Development

Contributions are welcome: bug fixes, tests, documentation, and features that stay within the library’s scope (attention and transformer building blocks on JAX / Flax).

If you are unsure whether something fits, open an issue first.

How to contribute

Use GitHub issues for bugs and design questions, and pull requests for code or doc changes. The usual workflow is: fork, branch, change, run the tests and doc build locally, open a PR.

What is expected of a PR:

• New behavior comes with tests under tests/.

• If the public API changes, the relevant docstrings and any affected pages in docs/ are updated in the same PR.

The project is distributed under the Apache License 2.0. There is no separate contributor agreement: if you open a PR that is merged, your contribution is included under the same license as the rest of the repository.

Local install

From a clone of the repository:

cd attnax
pip install -e ".[test]"

That installs Attnax in editable mode plus pytest and optax (used by the training tests and the getting-started notebook).

Tests

python -m pytest tests/

Writing docstrings

attnax uses Google-style docstrings rendered by sphinx.ext.napoleon. Inline markup inside a docstring is reStructuredText, not Markdown (page-level docs — getting_started.ipynb, this page, the README — are the other way around: Markdown / MyST, never RST).

A typical docstring for a new public function or class:

def my_function(x: jnp.ndarray, ...) -> jnp.ndarray:
  """One-line description ending with a period.

  Longer description if helpful, including mathematical context for
  algorithms.

  Args:
    x: Description of ``x``.
    ...

  Returns:
    Description of the return value.

  Raises:
    ValueError: When the inputs do not satisfy ...

  Example:
    ::

      import jax.numpy as jnp
      from attnax import my_function

      y = my_function(jnp.ones((2, 4)))
  """

Existing classes such as TransformerConfig and MultiHeadAttention are good references when in doubt.

Doctests

Code blocks inside Example: sections can be executed as doctests:

python -m doctest -v attnax/<module>.py

docs/conf.py configures doctest_global_setup so jax, jax.numpy as jnp, flax.nnx as nnx, and attnax are already imported in every doctest block — examples should reuse those names rather than re-importing.

Building the documentation

Install the docs extra (Sphinx, MyST-NB, the theme, etc.):

pip install -e ".[docs]"
cd docs
make html

Open docs/_build/html/index.html. The getting-started guide is a notebook (getting_started.ipynb); the Sphinx build renders it without re-executing cells (nb_execution_mode = "off" in conf.py).

Read the Docs uses docs/requirements.txt in addition to installing the package; keep that file aligned with the [docs] dependencies in pyproject.toml if you add doc-only packages.

Code style

The codebase follows the Google Python Style Guide. ruff enforces the rules wired in pyproject.toml.

AI-generated contributions

AI coding tools are fine to use; the bar for accepting a contribution is the same as for human-written code. In practice that means a contribution should not cost the reviewer more time than the contributor put into preparing it — read what you submit, run the tests, and check that the docstring template and code style match the rest of the package before opening the PR.