Host Autodiff Interop¶
Advect can wrap a NumPy-backed function as one differentiable operation inside PyTorch, JAX, or HIPS Autograd. The outer framework keeps its arrays and computation graph, while Advect supplies the wrapped function's VJP. The host-framework tutorial shows the pattern in a complete example.
Install and import¶
The base advect import loads no host framework. Install and import only the
bridge you use:
| Framework | Extra | Entry point | Reverse-mode execution |
|---|---|---|---|
| PyTorch | advect[torch] |
advect.interop.torch.wrap(function) |
Retains and consumes the forward Advect pullback |
| JAX | advect[jax] |
advect.interop.jax.wrap(function, ...) |
Executes eagerly or uses callbacks with a JIT/shape contract |
| HIPS Autograd | advect[autograd] |
advect.interop.autograd.wrap(function) |
Retains the reusable forward Advect linearization |
Shared contract¶
The callable accepts one or more positional or keyword tuple, list, or
dictionary pytrees, and every supplied leaf is differentiated. PyTorch leaves
are tensors, JAX leaves are arrays, and HIPS Autograd also accepts NumPy or
Python numeric scalars. Custom containers are supported only when both Advect
and the host framework recognize the same structure. Close over static
configuration rather than passing static leaves. Inputs and differentiable
outputs use standard NumPy floating or complex dtypes and outputs are nonempty
pytrees. JAX may additionally return a nondifferentiable auxiliary pytree with
has_aux=True.
All three bridges are first-order reverse-mode boundaries. They do not support Advect staging, host forward mode, or higher derivatives. The adapters handle the frameworks' different complex cotangent conventions, so native host losses over complex outputs receive the gradient convention expected by that host.
$ █