cottax.ode.consumers.execution module

Offers Diffrax, what runs an Integrator in jax, chosen per contract in ImplementDrivers.

class Diffrax(solver: AbstractSolver = Tsit5(), rtol: float = 1e-06, atol: float = 1e-06, max_steps: int = 4096, adjoint: AbstractAdjoint = RecursiveCheckpointAdjoint())

Bases: DriverImplementation[Integrator]

diffrax.diffeqsolve from the condition’s initial values over its span, an adaptive step from rtol and atol, saving the end, and in the same solve the Times where a Trajectory is asked back. The state is the unknowns as they are, a tuple, and the context its args. adjoint is how it differentiates: the default reverse mode only, diffrax.ForwardMode() forward mode only, diffrax.DirectAdjoint() both.

needs: AbstractClassVar[frozenset[type[Role]]] = frozenset({})

provides: AbstractClassVar[frozenset[type[Role]]] = frozenset({<class 'cottax.ode.contracts.Trajectory'>})

solver: AbstractSolver = Tsit5()

rtol: float = 1e-06

atol: float = 1e-06

max_steps: int = 4096

adjoint: AbstractAdjoint = RecursiveCheckpointAdjoint()

implements

alias of Integrator