eqiora.jax

← Python API index

Use Eqiora differentiable programs with JAX.

View source

Example: differentiate a scalar objective

Install the jax extra, then create program with the differentiation example:

python

import jax
import jax.numpy as jnp
from eqiora.jax import bind

jax.config.update("jax_enable_x64", True)
operator = bind(program)
parameters = (jnp.asarray(1.5, dtype=jnp.float64),)
gradient = jax.grad(lambda p: jnp.square(operator(p)[0]).sum())(parameters)
print(gradient)

This differentiates the sum of squared output values. Keep the parameter leaves together on one CPU device with their exact native shapes and dtypes. See framework adapters for JIT and batching examples.

API

JaxProgram

JAX view of a program in this process, with first-order vmap support.

python

class JaxProgram:
    def __init__(self, program: DifferentiableProgram) -> None: ...
    @property
    def program(self) -> DifferentiableProgram: ...
    @property
    def inputs(self) -> tuple[DifferentiableValueSignature, ...]: ...
    @property
    def outputs(self) -> tuple[DifferentiableValueSignature, ...]: ...
    def __call__(self, parameters: Sequence[jax.Array]) -> tuple[jax.Array, ...]: ...

bind

Wrap a DifferentiableProgram so that JAX can call and differentiate it.

python

def bind(program: DifferentiableProgram) -> JaxProgram: ...