onnx-ad 0.1.0__tar.gz

onnx_ad-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Joris Gillis
+
+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.

onnx_ad-0.1.0/PKG-INFO

@@ -0,0 +1,215 @@
+Metadata-Version: 2.4
+Name: onnx-ad
+Version: 0.1.0
+Summary: Source-code-transforming automatic differentiation on ONNX graphs
+Author: Joris Gillis
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/yacoda/onnx-ad
+Project-URL: Repository, https://github.com/yacoda/onnx-ad
+Project-URL: Issues, https://github.com/yacoda/onnx-ad/issues
+Keywords: onnx,automatic-differentiation,jacobian,adjoint,graph-transform
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: onnx>=1.14
+Requires-Dist: numpy>=1.21
+Provides-Extra: test
+Requires-Dist: onnxruntime>=1.16; extra == "test"
+Dynamic: license-file
+
+# onnx-ad
+
+Automatic differentiation **of ONNX graphs**, by source-code transformation: given a model,
+produce new ONNX models that compute its Jacobian-vector and vector-Jacobian products. Pure
+Python over the `onnx` protobuf, no runtime dependency, no framework in the loop.
+
+```sh
+python -m pip install onnx-ad
+```
+
+```python
+import onnx
+from onnx_ad import forward, reverse, family
+
+model = onnx.load("f.onnx")                      # inputs x -> outputs y
+onnx.save(forward(model), "fwd_f.onnx")          # + fwd_x -> + fwd_y = J . fwd_x
+onnx.save(reverse(model), "adj_f.onnx")          # + adj_y -> + adj_x = J^T . adj_y
+
+family(model, "generated/f.onnx")                # the whole set CasADi discovers
+```
+
+A derivative model keeps the original signature as a prefix and appends the seeds, so the
+primal outputs stay available. Any number of seed directions is evaluated in a single pass.
+
+## Why differentiate ONNX
+
+| Route | What it costs |
+| --- | --- |
+| Differentiate in PyTorch, then export | derivative graphs only for models that came from PyTorch; forward mode goes through `jvp`/`vmap`, which is where export breaks; every derivative order needs another trace |
+| Complex step (`Im f(x + i h v)/h`) | forward mode only, one evaluation per direction, and a convention rather than an identity at piecewise operations |
+| **This** | needs a rule per ONNX operation — but then any ONNX model has derivatives, from any producer, at any order |
+
+The third route is the one with no ceiling. A Jacobian-vector-product graph is itself an
+ONNX model, so it can be differentiated again, and consumers need no new capability.
+
+## Composition: second derivatives for free
+
+The passes keep the primal graph and emit only ordinary ONNX operations, so their own output
+is differentiable. Forward-over-adjoint — the exact-Hessian building block — is the two
+passes composed, with no special casing:
+
+```python
+adjoint = reverse(model)                 # x, adj_y -> adj_x
+hessian = forward(adjoint)               # + fwd_x, fwd_adj_y -> + fwd_adj_x
+```
+
+Repeated differentiation names itself the way CasADi's `diff_prefix` does: a model that
+already carries `fwd_x` gets `fwd2_`/`nfwd2` next, so `forward(forward(model))` needs no
+arguments.
+
+## CasADi
+
+The conventions are CasADi's, by default, so an emitted family drops into its ONNX backend
+with nothing to adapt:
+
+* **names** — `fwd_<x>`, `adj_<y>`, then `fwd2_`, `adj2_`, from the same rule
+  `FunctionInternal::diff_prefix` applies, and seed dimensions `nfwd`, `nadj`, `nfwd2`;
+* **layout** — CasADi reads an ONNX tensor as a matrix (rank 0/1 as a column, rank 2
+  directly, higher ranks flattened) and wants the seeds of an `r`-by-`c` value as one
+  `r`-by-`(nseed*c)` matrix. That is what `layout="casadi"` emits. Pass `layout="onnx"` for
+  the internal form instead, where the seed count is a trailing axis on the primal's own
+  shape;
+* **files** — `family` writes `f.onnx`, `adj_f.onnx` and `fwd_adj_f.onnx`, the
+  `<kind>_<filename>` siblings the backend looks for beside a model.
+
+```python
+f = casadi.GraphBuilder("generated/f.onnx").create("f")
+f.reverse(1)(x, f(x), w)                            # from adj_f.onnx
+casadi.hessian(casadi.dot(f(v), w), v)              # from fwd_adj_f.onnx
+```
+
+`examples/torch_to_casadi.py` exports a PyTorch model's **primal only** and generates the
+rest here; `examples/casadi_side.py` consumes it and checks gradients, Jacobians and exact
+Hessians against PyTorch. Verified against a CasADi build with `WITH_ONNX=ON` and
+`WITH_ONNX_RUNTIME=ON`, with `CASADI_ONNXRUNTIME_LIB` pointing at `libonnxruntime.so`.
+
+Two traps when the primal comes from `torch.onnx.export`: pass `external_data=False`, or the
+weights land in a sidecar `f.onnx.data` that CasADi cannot follow (it hands the model to ONNX
+Runtime as bytes); and CasADi needs the *complete* forward signature of the adjoint, which is
+why `family` seeds `adj_y` as well as `x`.
+
+Deliberately **not** offered: a `jacobian` pass. CasADi builds dense Jacobians from the
+adjoint itself, and a `jac_` sibling would only duplicate that.
+
+## What it emits
+
+Two walks over one rule table. Forward carries a tangent per value in graph order; reverse
+walks backwards, accumulating a contribution per value and summing where a value has several
+consumers. A value with no derivative is *absent* rather than zero, which is what keeps the
+emitted graph the size of the primal one — every weight in a network is such a value, and
+costs nothing.
+
+Nonlinear rules read the primal tensors they need straight from the primal graph rather than
+recomputing them: the tangent of `Tanh` is `(1 - y*y) * t`, with `y` the tensor the primal
+`Tanh` already produced. The reverse model therefore stays a plain function of `(x, adj_y)` —
+no "uses output" convention, nothing to wire up.
+
+Reverse mode's sharp edge is broadcasting: a contribution arrives shaped like the *result*
+and must be summed back over the axes the operand was broadcast along. Where the operand's
+shape is declared this is a static axis list; where it is symbolic the axes are computed at
+run time, so a dynamic batch dimension survives.
+
+### Operations with rules
+
+**Arithmetic** `Add` `Sub` `Mul` `Div` `Neg` `Pow` `Sum` `Mean` `Identity`
+
+**Elementwise** `Exp` `Log` `Sqrt` `Reciprocal` `Abs` `Sign` `Sin` `Cos` `Tan` `Sinh`
+`Cosh` `Asin` `Acos` `Atan` `Asinh` `Acosh` `Atanh` `Erf` `Tanh` `Sigmoid` `Relu`
+`LeakyRelu` `Elu` `Selu` `Celu` `PRelu` `ThresholdedRelu` `Softplus` `Softsign` `Shrink`
+`HardSigmoid` `HardSwish` `Mish` `Gelu` (exact and `tanh`)
+
+**Linear algebra** `MatMul` `Gemm` `Conv` (strided, dilated, grouped, 1-D and up; weight and
+bias too)
+
+**Shape** `Reshape` `Flatten` `Transpose` `Squeeze` `Unsqueeze` `Expand` `Concat` `Split`
+`Slice` `Pad` `Tile` `Gather` `GatherND` `CumSum`
+
+**Reductions** `ReduceSum` `ReduceMean` `ReduceMax` `ReduceMin` `ReduceProd`
+`ReduceLogSumExp` `ReduceL1` `ReduceL2` `ReduceSumSquare`
+
+**Selection** `Where` `Clip` `Min` `Max`
+
+**Networks** `Softmax` `LogSoftmax` `LayerNormalization` `BatchNormalization` (inference)
+`Dropout` (inference) `Cast` `CastLike`
+
+**Zero derivative, and allowed to consume differentiated values** `Shape` `Size` `NonZero`
+`Equal` `Greater` `Less` `GreaterOrEqual` `LessOrEqual` `And` `Or` `Xor` `Not` the bitwise
+family `ArgMax` `ArgMin` `IsNaN` `IsInf` `Floor` `Ceil` `Round` `Hardmax` `OneHot` `Det`
+
+An operation without a rule is an error **only when a differentiated value reaches it** — a
+`Resize` on a constant branch is fine.
+
+Still missing: pooling (`MaxPool`, `AveragePool`, `GlobalAveragePool`), `ConvTranspose` as a
+primal, `Einsum`, `Resize`, the scatter operations, `LpNormalization`,
+`InstanceNormalization` and `GroupNormalization`. Out of scope for a first release: control
+flow (`Loop`, `Scan`, `If`), sparsity, and training-mode operations.
+
+### Against PyTorch
+
+Every one of these exports (`torch.onnx.export`, `dynamo=True`, opset 18) differentiates in
+both modes, and the resulting Jacobian matches `torch.autograd.functional.jacobian` to
+float32 precision:
+
+| Model | Operations | forward | reverse |
+| --- | --- | ---: | ---: |
+| MLP | `Gemm` `Tanh` | 6e-8 | 6e-8 |
+| GELU MLP | `Gemm` `Erf` `Mul` `Div` `Add` | 2e-7 | 1e-7 |
+| SiLU MLP | `Gemm` `Sigmoid` `Mul` | 8e-8 | 8e-8 |
+| Attention block | `Gemm` `MatMul` `Softmax` `LayerNormalization` `Reshape` `Transpose` | 2e-7 | 2e-7 |
+| Convolutional net | `Conv` `Relu` `Gemm` `Reshape` | 1e-7 | 8e-8 |
+| Indexing and reductions | `GatherND` `Slice` `Clip` `ReduceMax` `ReduceProd` `ReduceL2` `ReduceLogSumExp` | 1e-7 | 1e-7 |
+| GRU cell | `Gemm` `Split` `Sigmoid` `Tanh` | 6e-8 | 6e-8 |
+
+## Running the result
+
+Reverse mode emits `Transpose` feeding `MatMul`, which ONNX Runtime's extended optimizer
+fuses into `com.microsoft.FusedMatMul` — a kernel registered for `float` only. On a
+double-precision model, load with the fusions off:
+
+```python
+options = ort.SessionOptions()
+options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
+session = ort.InferenceSession("adj_f.onnx", options)
+```
+
+## Command line
+
+```sh
+onnx-ad forward f.onnx fwd_f.onnx
+onnx-ad reverse f.onnx adj_f.onnx --inputs x --outputs y
+onnx-ad family  f.onnx generated/f.onnx
+```
+
+## Testing
+
+Every rule is executed through ONNX Runtime and compared against references that know nothing
+of the rule table: central finite differences of the primal model, an analytic Jacobian
+written in numpy, and forward against reverse — `J` and `J^T` come from separate walks over
+separate rules, so their agreement to machine precision is a real check.
+
+No PyTorch is involved in the unit tests; comparisons against it belong in an integration
+suite, not here.
+
+```sh
+python -m pip install -e ".[test]"
+python -m unittest discover -s tests -v
+```
+
+## License
+
+MIT. Releasing is documented in [RELEASE.md](RELEASE.md): PyPI Trusted Publishing, so no API
+token exists, and every artifact carries a PEP 740 provenance attestation.

onnx_ad-0.1.0/README.md

@@ -0,0 +1,192 @@
+# onnx-ad
+
+Automatic differentiation **of ONNX graphs**, by source-code transformation: given a model,
+produce new ONNX models that compute its Jacobian-vector and vector-Jacobian products. Pure
+Python over the `onnx` protobuf, no runtime dependency, no framework in the loop.
+
+```sh
+python -m pip install onnx-ad
+```
+
+```python
+import onnx
+from onnx_ad import forward, reverse, family
+
+model = onnx.load("f.onnx")                      # inputs x -> outputs y
+onnx.save(forward(model), "fwd_f.onnx")          # + fwd_x -> + fwd_y = J . fwd_x
+onnx.save(reverse(model), "adj_f.onnx")          # + adj_y -> + adj_x = J^T . adj_y
+
+family(model, "generated/f.onnx")                # the whole set CasADi discovers
+```
+
+A derivative model keeps the original signature as a prefix and appends the seeds, so the
+primal outputs stay available. Any number of seed directions is evaluated in a single pass.
+
+## Why differentiate ONNX
+
+| Route | What it costs |
+| --- | --- |
+| Differentiate in PyTorch, then export | derivative graphs only for models that came from PyTorch; forward mode goes through `jvp`/`vmap`, which is where export breaks; every derivative order needs another trace |
+| Complex step (`Im f(x + i h v)/h`) | forward mode only, one evaluation per direction, and a convention rather than an identity at piecewise operations |
+| **This** | needs a rule per ONNX operation — but then any ONNX model has derivatives, from any producer, at any order |
+
+The third route is the one with no ceiling. A Jacobian-vector-product graph is itself an
+ONNX model, so it can be differentiated again, and consumers need no new capability.
+
+## Composition: second derivatives for free
+
+The passes keep the primal graph and emit only ordinary ONNX operations, so their own output
+is differentiable. Forward-over-adjoint — the exact-Hessian building block — is the two
+passes composed, with no special casing:
+
+```python
+adjoint = reverse(model)                 # x, adj_y -> adj_x
+hessian = forward(adjoint)               # + fwd_x, fwd_adj_y -> + fwd_adj_x
+```
+
+Repeated differentiation names itself the way CasADi's `diff_prefix` does: a model that
+already carries `fwd_x` gets `fwd2_`/`nfwd2` next, so `forward(forward(model))` needs no
+arguments.
+
+## CasADi
+
+The conventions are CasADi's, by default, so an emitted family drops into its ONNX backend
+with nothing to adapt:
+
+* **names** — `fwd_<x>`, `adj_<y>`, then `fwd2_`, `adj2_`, from the same rule
+  `FunctionInternal::diff_prefix` applies, and seed dimensions `nfwd`, `nadj`, `nfwd2`;
+* **layout** — CasADi reads an ONNX tensor as a matrix (rank 0/1 as a column, rank 2
+  directly, higher ranks flattened) and wants the seeds of an `r`-by-`c` value as one
+  `r`-by-`(nseed*c)` matrix. That is what `layout="casadi"` emits. Pass `layout="onnx"` for
+  the internal form instead, where the seed count is a trailing axis on the primal's own
+  shape;
+* **files** — `family` writes `f.onnx`, `adj_f.onnx` and `fwd_adj_f.onnx`, the
+  `<kind>_<filename>` siblings the backend looks for beside a model.
+
+```python
+f = casadi.GraphBuilder("generated/f.onnx").create("f")
+f.reverse(1)(x, f(x), w)                            # from adj_f.onnx
+casadi.hessian(casadi.dot(f(v), w), v)              # from fwd_adj_f.onnx
+```
+
+`examples/torch_to_casadi.py` exports a PyTorch model's **primal only** and generates the
+rest here; `examples/casadi_side.py` consumes it and checks gradients, Jacobians and exact
+Hessians against PyTorch. Verified against a CasADi build with `WITH_ONNX=ON` and
+`WITH_ONNX_RUNTIME=ON`, with `CASADI_ONNXRUNTIME_LIB` pointing at `libonnxruntime.so`.
+
+Two traps when the primal comes from `torch.onnx.export`: pass `external_data=False`, or the
+weights land in a sidecar `f.onnx.data` that CasADi cannot follow (it hands the model to ONNX
+Runtime as bytes); and CasADi needs the *complete* forward signature of the adjoint, which is
+why `family` seeds `adj_y` as well as `x`.
+
+Deliberately **not** offered: a `jacobian` pass. CasADi builds dense Jacobians from the
+adjoint itself, and a `jac_` sibling would only duplicate that.
+
+## What it emits
+
+Two walks over one rule table. Forward carries a tangent per value in graph order; reverse
+walks backwards, accumulating a contribution per value and summing where a value has several
+consumers. A value with no derivative is *absent* rather than zero, which is what keeps the
+emitted graph the size of the primal one — every weight in a network is such a value, and
+costs nothing.
+
+Nonlinear rules read the primal tensors they need straight from the primal graph rather than
+recomputing them: the tangent of `Tanh` is `(1 - y*y) * t`, with `y` the tensor the primal
+`Tanh` already produced. The reverse model therefore stays a plain function of `(x, adj_y)` —
+no "uses output" convention, nothing to wire up.
+
+Reverse mode's sharp edge is broadcasting: a contribution arrives shaped like the *result*
+and must be summed back over the axes the operand was broadcast along. Where the operand's
+shape is declared this is a static axis list; where it is symbolic the axes are computed at
+run time, so a dynamic batch dimension survives.
+
+### Operations with rules
+
+**Arithmetic** `Add` `Sub` `Mul` `Div` `Neg` `Pow` `Sum` `Mean` `Identity`
+
+**Elementwise** `Exp` `Log` `Sqrt` `Reciprocal` `Abs` `Sign` `Sin` `Cos` `Tan` `Sinh`
+`Cosh` `Asin` `Acos` `Atan` `Asinh` `Acosh` `Atanh` `Erf` `Tanh` `Sigmoid` `Relu`
+`LeakyRelu` `Elu` `Selu` `Celu` `PRelu` `ThresholdedRelu` `Softplus` `Softsign` `Shrink`
+`HardSigmoid` `HardSwish` `Mish` `Gelu` (exact and `tanh`)
+
+**Linear algebra** `MatMul` `Gemm` `Conv` (strided, dilated, grouped, 1-D and up; weight and
+bias too)
+
+**Shape** `Reshape` `Flatten` `Transpose` `Squeeze` `Unsqueeze` `Expand` `Concat` `Split`
+`Slice` `Pad` `Tile` `Gather` `GatherND` `CumSum`
+
+**Reductions** `ReduceSum` `ReduceMean` `ReduceMax` `ReduceMin` `ReduceProd`
+`ReduceLogSumExp` `ReduceL1` `ReduceL2` `ReduceSumSquare`
+
+**Selection** `Where` `Clip` `Min` `Max`
+
+**Networks** `Softmax` `LogSoftmax` `LayerNormalization` `BatchNormalization` (inference)
+`Dropout` (inference) `Cast` `CastLike`
+
+**Zero derivative, and allowed to consume differentiated values** `Shape` `Size` `NonZero`
+`Equal` `Greater` `Less` `GreaterOrEqual` `LessOrEqual` `And` `Or` `Xor` `Not` the bitwise
+family `ArgMax` `ArgMin` `IsNaN` `IsInf` `Floor` `Ceil` `Round` `Hardmax` `OneHot` `Det`
+
+An operation without a rule is an error **only when a differentiated value reaches it** — a
+`Resize` on a constant branch is fine.
+
+Still missing: pooling (`MaxPool`, `AveragePool`, `GlobalAveragePool`), `ConvTranspose` as a
+primal, `Einsum`, `Resize`, the scatter operations, `LpNormalization`,
+`InstanceNormalization` and `GroupNormalization`. Out of scope for a first release: control
+flow (`Loop`, `Scan`, `If`), sparsity, and training-mode operations.
+
+### Against PyTorch
+
+Every one of these exports (`torch.onnx.export`, `dynamo=True`, opset 18) differentiates in
+both modes, and the resulting Jacobian matches `torch.autograd.functional.jacobian` to
+float32 precision:
+
+| Model | Operations | forward | reverse |
+| --- | --- | ---: | ---: |
+| MLP | `Gemm` `Tanh` | 6e-8 | 6e-8 |
+| GELU MLP | `Gemm` `Erf` `Mul` `Div` `Add` | 2e-7 | 1e-7 |
+| SiLU MLP | `Gemm` `Sigmoid` `Mul` | 8e-8 | 8e-8 |
+| Attention block | `Gemm` `MatMul` `Softmax` `LayerNormalization` `Reshape` `Transpose` | 2e-7 | 2e-7 |
+| Convolutional net | `Conv` `Relu` `Gemm` `Reshape` | 1e-7 | 8e-8 |
+| Indexing and reductions | `GatherND` `Slice` `Clip` `ReduceMax` `ReduceProd` `ReduceL2` `ReduceLogSumExp` | 1e-7 | 1e-7 |
+| GRU cell | `Gemm` `Split` `Sigmoid` `Tanh` | 6e-8 | 6e-8 |
+
+## Running the result
+
+Reverse mode emits `Transpose` feeding `MatMul`, which ONNX Runtime's extended optimizer
+fuses into `com.microsoft.FusedMatMul` — a kernel registered for `float` only. On a
+double-precision model, load with the fusions off:
+
+```python
+options = ort.SessionOptions()
+options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
+session = ort.InferenceSession("adj_f.onnx", options)
+```
+
+## Command line
+
+```sh
+onnx-ad forward f.onnx fwd_f.onnx
+onnx-ad reverse f.onnx adj_f.onnx --inputs x --outputs y
+onnx-ad family  f.onnx generated/f.onnx
+```
+
+## Testing
+
+Every rule is executed through ONNX Runtime and compared against references that know nothing
+of the rule table: central finite differences of the primal model, an analytic Jacobian
+written in numpy, and forward against reverse — `J` and `J^T` come from separate walks over
+separate rules, so their agreement to machine precision is a real check.
+
+No PyTorch is involved in the unit tests; comparisons against it belong in an integration
+suite, not here.
+
+```sh
+python -m pip install -e ".[test]"
+python -m unittest discover -s tests -v
+```
+
+## License
+
+MIT. Releasing is documented in [RELEASE.md](RELEASE.md): PyPI Trusted Publishing, so no API
+token exists, and every artifact carries a PEP 740 provenance attestation.

onnx_ad-0.1.0/pyproject.toml

@@ -0,0 +1,35 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "onnx-ad"
+version = "0.1.0"
+description = "Source-code-transforming automatic differentiation on ONNX graphs"
+readme = "README.md"
+requires-python = ">=3.9"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{name = "Joris Gillis"}]
+keywords = ["onnx", "automatic-differentiation", "jacobian", "adjoint", "graph-transform"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3 :: Only",
+  "Topic :: Scientific/Engineering :: Mathematics",
+]
+dependencies = ["onnx>=1.14", "numpy>=1.21"]
+
+[project.optional-dependencies]
+test = ["onnxruntime>=1.16"]
+
+[project.scripts]
+onnx-ad = "onnx_ad.__main__:main"
+
+[project.urls]
+Homepage = "https://github.com/yacoda/onnx-ad"
+Repository = "https://github.com/yacoda/onnx-ad"
+Issues = "https://github.com/yacoda/onnx-ad/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]

onnx_ad-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

onnx_ad-0.1.0/src/onnx_ad/__init__.py

@@ -0,0 +1,20 @@
+"""Source-code-transforming automatic differentiation on ONNX graphs.
+
+`forward` and `reverse` rewrite an ONNX model into another ONNX model that computes its
+Jacobian-vector or vector-Jacobian products, by applying a rule per operation over the
+protobuf. No framework is in the loop: any ONNX model has derivatives, from any producer,
+and since a derivative model is itself an ONNX model, `forward(reverse(model))` gives
+forward-over-adjoint -- the exact-Hessian building block -- with no special casing.
+
+Several seed directions ride in one evaluation, and the conventions are CasADi's own: the
+derivative prefixes and seed dimensions its `diff_prefix` rule would pick (`fwd_`/`nfwd`,
+then `fwd2_`/`nfwd2`), the seed layout its 2-D reading of an ONNX tensor expects, and the
+`<kind>_<filename>` sibling names its ONNX backend discovers. `family` writes the whole set.
+"""
+from ._build import UnsupportedOperator
+from .family import family, sibling
+from .forward import forward
+from .reverse import reverse
+
+__all__ = ["forward", "reverse", "family", "sibling", "UnsupportedOperator", "__version__"]
+__version__ = "0.1.0"

onnx_ad-0.1.0/src/onnx_ad/__main__.py

@@ -0,0 +1,52 @@
+"""Command line: write the derivative model, or the whole sibling family, of an ONNX model."""
+import argparse
+
+import onnx
+
+from . import __version__, family, forward, reverse
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="onnx-ad", description=__doc__)
+    parser.add_argument("--version", action="version", version=__version__)
+    sub = parser.add_subparsers(dest="command", required=True)
+    for name, help_text in [
+            ("forward", "emit a model computing fwd_y = J . fwd_x"),
+            ("reverse", "emit a model computing adj_x = J^T . adj_y"),
+            ("family", "write the primal and the sibling derivative models CasADi discovers")]:
+        command = sub.add_parser(name, help=help_text)
+        command.add_argument("input")
+        command.add_argument("output")
+        command.add_argument("--inputs", nargs="+",
+                             help="graph inputs to differentiate [every floating-point one]")
+        command.add_argument("--outputs", nargs="+",
+                             help="graph outputs to differentiate [every floating-point one]")
+        if name == "family":
+            command.add_argument("--forward-sibling", action="store_true",
+                                 help="also write fwd_<name>, worth it when inputs < outputs")
+            command.add_argument("--no-second-order", dest="second_order",
+                                 action="store_false",
+                                 help="skip fwd_adj_<name>, the exact-Hessian sibling")
+        else:
+            command.add_argument("--prefix", help="derivative-tensor prefix [CasADi's rule]")
+            command.add_argument("--dim", help="symbolic seed dimension [CasADi's rule]")
+            command.add_argument("--layout", choices=["casadi", "onnx"], default="casadi",
+                                 help="seed layout: CasADi's packed matrix, or a trailing "
+                                      "axis on the primal's own shape [casadi]")
+    args = parser.parse_args(argv)
+    model = onnx.load(args.input)
+    if args.command == "family":
+        for path in family(model, args.output, forward_sibling=args.forward_sibling,
+                           second_order=args.second_order, inputs=args.inputs,
+                           outputs=args.outputs):
+            print(path)
+        return
+    pass_ = forward if args.command == "forward" else reverse
+    result = pass_(model, inputs=args.inputs, outputs=args.outputs, prefix=args.prefix,
+                   dim=args.dim, layout=args.layout)
+    onnx.checker.check_model(result)
+    onnx.save(result, args.output)
+
+
+if __name__ == "__main__":
+    main()