asdex 0.1.0__tar.gz

asdex-0.1.0/.claude/skills/add-handler/SKILL.md

@@ -0,0 +1,109 @@
+---
+name: add-handler
+description: Add a precise primitive handler to the jaxpr interpreter, replacing a conservative fallback.
+argument-hint: "[primitive-name]"
+disable-model-invocation: true
+---
+
+# Add precise handler for `$ARGUMENTS`
+
+Add a precise propagation handler for the `$ARGUMENTS` primitive,
+replacing the conservative fallback.
+
+## Workflow
+
+### 1. Research
+
+Do all research in this step using extended **planning mode**.
+
+Before writing code:
+- Read the JAX docs for the primitive: fetch `https://docs.jax.dev/en/latest/_autosummary/jax.lax.$ARGUMENTS.html`
+- Read `src/asdex/_interpret/CLAUDE.md` for conventions (docstring style, semantic line breaks, handler structure)
+- Read `src/asdex/_interpret/_commons.py` to understand available utilities
+- Read an existing handler with similar structure (e.g. `_pad.py`, `_transpose.py`, `_reduction.py`) as a reference
+- Read the existing test in `tests/_interpret/test_internals.py` for the primitive (search for `$ARGUMENTS`)
+- Read `src/asdex/_interpret/__init__.py` to see the current dispatch and fallback setup
+
+Understand the primitive's semantics:
+how do input and output element indices map to each other?
+What is the Jacobian structure (permutation, selection, block-diagonal, etc.)?
+
+### 2. Implement handler
+
+- Create `src/asdex/_interpret/_$ARGUMENTS.py` with `prop_$ARGUMENTS(eqn, deps)`.
+- Follow the handler docstring style from `_interpret/CLAUDE.md`.
+- If the primitive preserves or predictably transforms input values,
+  propagate `const_vals` so downstream handlers can stay precise.
+
+### 3. Wire up dispatch
+
+In `src/asdex/_interpret/__init__.py`:
+- Import the new handler
+- Add a `case "$ARGUMENTS":` branch in `prop_dispatch` calling the handler
+- Remove `"$ARGUMENTS"` from the conservative fallback `case` group
+
+### 4. Update tests
+
+In `tests/_interpret/test_internals.py`:
+- Update the existing test: change expected values from dense (`np.ones`) to the precise pattern
+- Remove the `@pytest.mark.fallback` marker and `TODO` comments
+
+Create `tests/_interpret/test_$ARGUMENTS.py` with thorough tests:
+- Multiple dimensionalities (1D, 2D, 3D, 4D where applicable)
+- Broadcasting shapes: size-1 dimensions that broadcast (e.g. `(3,4)` op `(3,1)`)
+- Non-square shapes (e.g. `(3,4)` not `(4,4)`) so dimension mix-ups are caught
+- Edge cases (size-0 dimensions, identity/trivial parameters)
+- Real-world usage patterns (e.g. `jnp` functions that lower to this primitive)
+- For at least one test per dimensionality, verify precision by comparing the detected
+  pattern against `(np.abs(jax.jacobian(f)(x)) > 1e-10)` using `assert_array_equal`.
+  Choose test functions that avoid local sparsity (e.g. multiply by zero) so the
+  numerical Jacobian matches the structural pattern.
+
+### 5. Verify
+
+Run in order:
+```bash
+uv run ruff check src/asdex/_interpret/_$ARGUMENTS.py
+uv run pytest tests/_interpret/test_$ARGUMENTS.py -v
+uv run pytest tests/_interpret/test_internals.py -v
+uv run pytest tests/ -x
+```
+
+### 6. Adversarial tests
+
+Reread the JAX docs for the primitive: fetch `https://docs.jax.dev/en/latest/_autosummary/jax.lax.$ARGUMENTS.html`
+
+Try to break the implementation by testing inputs the handler might not expect:
+
+- **Dimensionality**: 1D, 2D, 3D, and higher — if any are missing, add them.
+- **Asymmetric shapes**: always use non-square shapes (e.g. `(3,4)` not `(4,4)`) so that dimension transposition bugs are caught.
+- **Degenerate shapes**: size-0 dimensions, size-1 dimensions, scalar inputs (where the primitive supports them).
+- **Boundary parameters**: empty parameter lists, all-dimensions, single-dimension, negative indices (if applicable).
+- **Compositions**: the primitive chained with itself (e.g. double-reverse, transpose-of-transpose) or with related ops.
+- **Non-contiguous patterns**: inputs where dependencies are not simply `{i}` per element (e.g. from a prior broadcast or reduction) to verify `.copy()` and set merging behave correctly.
+- **Conservative audit**: for each test case, verify the result is strictly sparser than what `conservative_indices()` would produce. If the handler silently falls back on any shape the primitive supports, investigate.
+  If the fallback cannot be fixed immediately, you **must** add a `@pytest.mark.fallback` test with a `TODO(primitive)` comment showing the precise expected pattern.
+  Catching conservative patterns is extremely valuable for future development.
+- **Const chain**: if the primitive can appear between a literal and a downstream consumer (e.g. type conversions, reshapes, broadcasts on index arrays), write a test composing it with a downstream gather and verify the gather resolves precisely.
+
+For each new test, verify the expected output by hand or against `jax.jacobian`.
+Update and re-verify the handler if any test reveals a bug.
+
+### 7. Simplify
+
+Review the implementation with fresh eyes and look for opportunities to reduce complexity:
+
+- **Vectorize loops**: can per-element Python loops be replaced with numpy operations?
+  Pattern: build a flat permutation or index array with `np.arange`, `np.flip`, `np.transpose`, `np.indices`, or `np.ravel_multi_index`,
+  then index into `in_indices` in a single list comprehension.
+  See `_rev.py`, `_reshape.py`, `_concatenate.py`, and `_broadcast.py` for examples.
+- **Remove unused imports**: after vectorizing, utilities like `flat_to_coords`, `row_strides`, and `numel` may no longer be needed.
+- **Eliminate intermediate variables**: if a value is computed and used only once, inline it.
+- **Simplify special cases**: can a special-case branch be absorbed into the general case?
+
+After any change, re-run verification (step 6).
+
+### 8. Update docs
+
+- `TODO.md`: check off the primitive and its test items
+- `src/asdex/_interpret/CLAUDE.md`: add the new module to the file listing

asdex-0.1.0/.github/workflows/benchmarks.yml

@@ -0,0 +1,33 @@
+name: Benchmarks
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  Benchmarks:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+      deployments: write
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+      - name: Set up Python
+        run: uv python install 3.12
+      - name: Install dependencies
+        run: uv sync --group dev
+      - name: Run benchmarks
+        run: uv run pytest tests/test_benchmarks.py -m dashboard --benchmark-json=benchmark.json
+      - name: Store benchmark result
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          tool: 'pytest'
+          output-file-path: benchmark.json
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          auto-push: true
+          alert-threshold: '150%'
+          comment-on-alert: true
+          fail-on-alert: false

asdex-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,49 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  Tests:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.11", "3.14"]
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+      - name: Set up Python
+        run: uv python install ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: uv sync --group dev
+      - name: Test with coverage
+        run: uv run pytest -m "" --cov=src/asdex --cov-report=xml  # -m "" overrides default marker filter to include slow and benchmark tests
+      - name: Upload coverage to Codecov
+        if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.14'
+        uses: codecov/codecov-action@v5
+        with:
+          files: coverage.xml
+          token: ${{ secrets.CODECOV_TOKEN }}
+
+  Linting:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+      - name: Set up Python
+        run: uv python install 3.14
+      - name: Install dependencies
+        run: uv sync --group dev
+      - name: Lint
+        run: uv run ruff check .
+      - name: Format
+        run: uv run ruff format --check .
+      - name: Type check
+        run: uv run ty check

asdex-0.1.0/.github/workflows/docs.yml

@@ -0,0 +1,27 @@
+name: Documentation
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  docs:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - run: uv python install 3.12
+      - run: uv sync --group docs
+      - run: uv run mkdocs build --strict
+
+      - name: Deploy to GitHub Pages
+        if: github.event_name == 'push' && github.ref == 'refs/heads/main'
+        uses: peaceiris/actions-gh-pages@v4
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          publish_dir: ./site
+          keep_files: true

asdex-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,33 @@
+name: Release
+# Based on https://github.com/astral-sh/trusted-publishing-examples/blob/main/.github/workflows/release.yml
+
+on:
+  push:
+    tags:
+      # Publish on any tag starting with a `v`, e.g., v0.1.0
+      - v*
+
+jobs:
+  run:
+    name: Publish to PyPI
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+    permissions:
+      id-token: write
+      contents: read
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+      - name: Install Python 3.14
+        run: uv python install 3.14
+      - name: Build
+        run: uv build
+      - name: Install dependencies
+        run: uv sync --group dev
+      - name: Test
+        run: uv run pytest -m ""
+      - name: Publish
+        run: uv publish

asdex-0.1.0/.gitignore

@@ -0,0 +1,20 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+uv.lock
+.python-version
+.benchmarks
+.pytest_cache
+.ruff_cache
+
+# MkDocs build output
+site/
+
+# Virtual environments
+.venv
+.claude/settings.local.json
+_context

asdex-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,19 @@
+repos:
+  - repo: local
+    hooks:
+      - id: ruff
+        name: ruff
+        entry: uv run ruff check --fix
+        language: system
+        types: [python]
+      - id: ruff-format
+        name: ruff-format
+        entry: uv run ruff format
+        language: system
+        types: [python]
+      - id: ty
+        name: ty
+        entry: uv run ty check
+        language: system
+        pass_filenames: false
+        types: [python]

asdex-0.1.0/CLAUDE.md

@@ -0,0 +1,103 @@
+# asdex - Automatic Sparse Differentiation in JAX
+
+This package implements [Automatic Sparse Differentiation](https://iclr-blogposts.github.io/2025/blog/sparse-autodiff/) (ASD) in JAX.
+
+## Overview
+
+ASD exploits sparsity to reduce the cost of computing sparse Jacobians and Hessians:
+
+1. **Detection**: Analyze the jaxpr computation graph to detect the global sparsity pattern
+2. **Coloring**: Assign colors to rows so that rows sharing non-zero columns get different colors
+3. **Decompression**: Compute one VJP/HVP per color instead of one per row, then extract the sparse matrix
+
+## Structure
+
+```
+src/asdex/
+├── __init__.py         # Public API
+├── pattern.py          # SparsityPattern and ColoredPattern data structures
+├── detection.py        # Jacobian and Hessian sparsity detection via jaxpr analysis
+├── coloring.py         # Graph coloring (row, column, symmetric) and convenience functions
+├── decompression.py    # Sparse Jacobian (VJP/JVP) and Hessian (HVP) computation
+├── verify.py           # Correctness checks (check_jacobian_correctness, check_hessian_correctness)
+├── _display.py         # Display/formatting utilities
+└── _interpret/         # Custom jaxpr interpreter for index set propagation
+```
+
+The interpreter internals are described in `src/asdex/_interpret/CLAUDE.md`.
+The structure of the test folder is described in `tests/CLAUDE.md`.
+
+## Development
+
+```bash
+uv run ruff check --fix .  # lint + auto-fix
+uv run ruff format .       # format
+uv run ty check            # type check
+uv run pytest              # run tests
+```
+
+## Code style
+
+- Favor `match` statements over long if-else chains.
+  Use explicit cases and default to `case _ as unreachable: assert_never(unreachable)`.
+- Use plain `# Section name` comments for section separators, 
+  not banner-style `# -- Section name ---`.
+
+## Architecture
+
+### Jacobians
+
+```
+jacobian(f, input_shape)(x)          # one-call API
+jacobian_from_coloring(f, coloring)(x)  # from pre-computed coloring
+  │
+  ├─ 1. DETECTION
+  │     jacobian_sparsity(f, input_shape)
+  │     ├─ make_jaxpr(f) → jaxpr
+  │     ├─ prop_jaxpr() → index sets
+  │     └─ SparsityPattern
+  │
+  ├─ 2. COLORING
+  │     jacobian_coloring_from_sparsity(sparsity)
+  │
+  └─ 3. DECOMPRESSION
+        One VJP or JVP per color
+
+Precompute: jacobian_coloring(f, shape) = detect + color
+```
+
+### Hessians
+
+```
+hessian(f, input_shape)(x)          # one-call API
+hessian_from_coloring(f, coloring)(x)  # from pre-computed coloring
+  │
+  ├─ 1. DETECTION
+  │     hessian_sparsity(f, input_shape)
+  │     └─ jacobian_sparsity(grad(f), input_shape)
+  │
+  ├─ 2. COLORING
+  │     hessian_coloring_from_sparsity(sparsity)
+  │
+  └─ 3. DECOMPRESSION
+        One HVP per color (fwd-over-rev)
+
+Precompute: hessian_coloring(f, shape) = detect + color_symmetric
+```
+
+## Commits
+
+Use [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) for all commit messages (e.g. `feat:`, `fix:`, `docs:`, `refactor:`, `test:`).
+For breaking changes, add `!` after the type (e.g. `feat!:`).
+
+## Design philosophy
+
+When writing new code, adhere to these design principles:
+
+- **Minimize complexity**: The primary goal of software design is to minimize complexity—anything that makes a system hard to understand and modify.
+
+- **Information hiding**: Each module should encapsulate design decisions that other modules don't need to know about, preventing information leakage across boundaries.
+
+- **Pull complexity downward**: It's better for a module to be internally complex if it keeps the interface simple for others. Don't expose complexity to callers.
+
+- **Favor exceptions over wrong results**: Raise errors for unknown edge cases rather than guessing. 

asdex-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Adrian Hill <gh@adrianhill.de>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

asdex-0.1.0/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.4
+Name: asdex
+Version: 0.1.0
+Summary: Automatic Sparse Differentiation in JAX
+Project-URL: Homepage, https://github.com/adrhill/asdex
+Project-URL: Documentation, https://adrianhill.de/asdex
+Project-URL: Benchmarks, https://adrianhill.de/asdex/benchmarks
+Project-URL: Repository, https://github.com/adrhill/asdex
+Project-URL: Issues, https://github.com/adrhill/asdex/issues
+Author-email: Adrian Hill <gh@adrianhill.de>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: automatic-differentiation,automatic-sparse-differentiation,graph-coloring,hessian,jacobian,jax,jaxpr,scientific-computing,sparse-automatic-differentiation,sparsity,sparsity-detection
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.9.0
+Requires-Dist: numpy>=1.26.0
+Description-Content-Type: text/markdown
+
+# asdex
+
+[![CI](https://github.com/adrhill/asdex/actions/workflows/ci.yml/badge.svg)](https://github.com/adrhill/asdex/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/adrhill/asdex/graph/badge.svg)](https://codecov.io/gh/adrhill/asdex)
+[![Docs](https://img.shields.io/badge/docs-online-blue)](https://adrianhill.de/asdex/)
+[![Benchmarks](https://img.shields.io/badge/benchmarks-view-blue)](https://adrianhill.de/asdex/dev/bench/)
+
+[Automatic Sparse Differentiation](https://iclr-blogposts.github.io/2025/blog/sparse-autodiff/) in JAX.
+
+`asdex` (pronounced _Aztecs_) exploits sparsity structure to efficiently compute sparse Jacobians and Hessians.
+It implements a custom [Jaxpr](https://docs.jax.dev/en/latest/jaxpr.html) interpreter
+that uses [abstract interpretation](https://en.wikipedia.org/wiki/Abstract_interpretation)
+to detect sparsity patterns from the computation graph,
+then uses graph coloring to minimize the number of AD passes needed.
+
+> [!WARNING]
+> `asdex` is in early development.
+> The API may change without notice.
+> Use at your own risk.
+
+## Installation
+
+```bash
+pip install git+https://github.com/adrhill/asdex.git
+```
+
+Or with [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv add git+https://github.com/adrhill/asdex.git
+```
+
+## Example
+
+```python
+import numpy as np
+from asdex import jacobian
+
+def f(x):
+    return (x[1:] - x[:-1]) ** 2
+
+jac_fn = jacobian(f, input_shape=50)
+# ColoredPattern(49×50, nnz=98, sparsity=96.0%, JVP, 2 colors)
+#   2 JVPs (instead of 49 VJPs or 50 JVPs)
+# ⎡⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎤   ⎡⣿⎤
+# ⎢⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥ → ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⎥   ⎢⣿⎥
+# ⎣⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⎦   ⎣⠉⎦
+
+for x in inputs:
+    J = jac_fn(x)
+```
+
+Instead of 49 VJPs or 50 JVPs,
+`asdex` computes the full sparse Jacobian with just 2 JVPs.
+
+## Documentation
+
+- [Getting Started](https://adrianhill.de/asdex/tutorials/getting-started/) — step-by-step tutorial
+- [How-To Guides](https://adrianhill.de/asdex/how-to/jacobians/) — task-oriented recipes
+- [Explanation](https://adrianhill.de/asdex/explanation/sparsity-detection/) — how and why it works
+- [API Reference](https://adrianhill.de/asdex/reference/) — full API documentation
+
+## Acknowledgements
+
+This package is built with Claude Code based on previous work by [Adrian Hill](https://github.com/adrhill), [Guillaume Dalle](https://github.com/gdalle), and [Alexis Montoison](https://github.com/amontoison) in the [Julia programming language](https://julialang.org):
+
+- [_An Illustrated Guide to Automatic Sparse Differentiation_](https://iclr-blogposts.github.io/2025/blog/sparse-autodiff/), A. Hill, G. Dalle, A. Montoison (2025)
+- [_Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians_](https://openreview.net/forum?id=GtXSN52nIW), A. Hill & G. Dalle (2025)
+- [_Revisiting Sparse Matrix Coloring and Bicoloring_](https://arxiv.org/abs/2505.07308), A. Montoison, G. Dalle, A. Gebremedhin (2025)
+- [_SparseConnectivityTracer.jl_](https://github.com/adrhill/SparseConnectivityTracer.jl), A. Hill, G. Dalle
+- [_SparseMatrixColorings.jl_](https://github.com/gdalle/SparseMatrixColorings.jl), G. Dalle, A. Montoison
+- [_sparsediffax_](https://github.com/gdalle/sparsediffax), G. Dalle
+
+which in turn stands on the shoulders of giants — notably Andreas Griewank, Andrea Walther, and Assefaw Gebremedhin.

asdex-0.1.0/README.md

@@ -0,0 +1,84 @@
+# asdex
+
+[![CI](https://github.com/adrhill/asdex/actions/workflows/ci.yml/badge.svg)](https://github.com/adrhill/asdex/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/adrhill/asdex/graph/badge.svg)](https://codecov.io/gh/adrhill/asdex)
+[![Docs](https://img.shields.io/badge/docs-online-blue)](https://adrianhill.de/asdex/)
+[![Benchmarks](https://img.shields.io/badge/benchmarks-view-blue)](https://adrianhill.de/asdex/dev/bench/)
+
+[Automatic Sparse Differentiation](https://iclr-blogposts.github.io/2025/blog/sparse-autodiff/) in JAX.
+
+`asdex` (pronounced _Aztecs_) exploits sparsity structure to efficiently compute sparse Jacobians and Hessians.
+It implements a custom [Jaxpr](https://docs.jax.dev/en/latest/jaxpr.html) interpreter
+that uses [abstract interpretation](https://en.wikipedia.org/wiki/Abstract_interpretation)
+to detect sparsity patterns from the computation graph,
+then uses graph coloring to minimize the number of AD passes needed.
+
+> [!WARNING]
+> `asdex` is in early development.
+> The API may change without notice.
+> Use at your own risk.
+
+## Installation
+
+```bash
+pip install git+https://github.com/adrhill/asdex.git
+```
+
+Or with [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv add git+https://github.com/adrhill/asdex.git
+```
+
+## Example
+
+```python
+import numpy as np
+from asdex import jacobian
+
+def f(x):
+    return (x[1:] - x[:-1]) ** 2
+
+jac_fn = jacobian(f, input_shape=50)
+# ColoredPattern(49×50, nnz=98, sparsity=96.0%, JVP, 2 colors)
+#   2 JVPs (instead of 49 VJPs or 50 JVPs)
+# ⎡⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎤   ⎡⣿⎤
+# ⎢⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥ → ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⎥   ⎢⣿⎥
+# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⎥   ⎢⣿⎥
+# ⎣⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⎦   ⎣⠉⎦
+
+for x in inputs:
+    J = jac_fn(x)
+```
+
+Instead of 49 VJPs or 50 JVPs,
+`asdex` computes the full sparse Jacobian with just 2 JVPs.
+
+## Documentation
+
+- [Getting Started](https://adrianhill.de/asdex/tutorials/getting-started/) — step-by-step tutorial
+- [How-To Guides](https://adrianhill.de/asdex/how-to/jacobians/) — task-oriented recipes
+- [Explanation](https://adrianhill.de/asdex/explanation/sparsity-detection/) — how and why it works
+- [API Reference](https://adrianhill.de/asdex/reference/) — full API documentation
+
+## Acknowledgements
+
+This package is built with Claude Code based on previous work by [Adrian Hill](https://github.com/adrhill), [Guillaume Dalle](https://github.com/gdalle), and [Alexis Montoison](https://github.com/amontoison) in the [Julia programming language](https://julialang.org):
+
+- [_An Illustrated Guide to Automatic Sparse Differentiation_](https://iclr-blogposts.github.io/2025/blog/sparse-autodiff/), A. Hill, G. Dalle, A. Montoison (2025)
+- [_Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians_](https://openreview.net/forum?id=GtXSN52nIW), A. Hill & G. Dalle (2025)
+- [_Revisiting Sparse Matrix Coloring and Bicoloring_](https://arxiv.org/abs/2505.07308), A. Montoison, G. Dalle, A. Gebremedhin (2025)
+- [_SparseConnectivityTracer.jl_](https://github.com/adrhill/SparseConnectivityTracer.jl), A. Hill, G. Dalle
+- [_SparseMatrixColorings.jl_](https://github.com/gdalle/SparseMatrixColorings.jl), G. Dalle, A. Montoison
+- [_sparsediffax_](https://github.com/gdalle/sparsediffax), G. Dalle
+
+which in turn stands on the shoulders of giants — notably Andreas Griewank, Andrea Walther, and Assefaw Gebremedhin.