Contributing¶

We warmly welcome contributions!

How You Can Help¶

Add a plugin: Extend jax2onnx by writing a simple Python file in jax2onnx/plugins: a primitive or an example. The Plugin Quickstart walks through the process step-by-step.
Bug fixes & improvements: PRs and issues are always welcome on GitHub.

Getting Started¶

Fork the repository and clone it locally.
Install dependencies with Poetry:
```
poetry install
```
Install pre-commit hooks (essential for linting):
```
poetry run pre-commit install
```

Development Workflow¶

Linting & Formatting¶

We use ruff (linting), black (formatting), and mypy via scripts/check_typing.sh.

# Run all checks (recommended)
poetry run pre-commit run --all-files

# Or run individually
poetry run ruff check .
poetry run black .
./scripts/check_typing.sh

Running Tests¶

# Run all tests
poetry run pytest

# Run a specific test file
poetry run pytest tests/extra_tests/framework/test_no_onnx_in_converter_plugins.py

Repository Helper Scripts¶

# Run pre-commit, regenerate tests, then run pytest (fails fast on first error)
./scripts/run_all_checks.sh

# Optional: include MaxText SotA checks in the same runner
JAX2ONNX_RUN_MAXTEXT=1 ./scripts/run_all_checks.sh

# Refresh autogenerated coverage/examples pages
# This also regenerates tests and runs pytest to collect current statuses
poetry run python scripts/generate_readme.py

Documentation¶

To preview documentation changes locally:

poetry run mkdocs serve
poetry run mkdocs build --strict

See the Plugin System documentation for details on adding new operators.