Skip to content

treams_rs.testing

Finite-difference checks of first-order gradients, pullbacks and pushforwards.

A pullback maps the gradient of a real loss with respect to an output to the gradients with respect to the inputs; treams_rs.diff defines records, contexts and pullbacks. Complex gradients follow dL = Re(vdot(g, dx)). These checks are test tools: treams-rs never computes a gradient by finite differences.

check_gradient

check_gradient(function: Callable[..., ArrayLike], gradient: Callable[..., ArrayLike | tuple[ArrayLike, ...]], *parameters: ArrayLike, directions: tuple[ArrayLike, ...] | None=None, step: float=1e-06, rtol: float=1e-05, atol: float=1e-07, seed: int=0) -> None

Check the gradient of a real scalar objective against finite differences.

The check calls gradient once and function 1 + 2n times for n parameters. It runs check_pullback on the record (function(*p), lambda g: gradient(*p)) with the cotangent 1.0.

Args

  • function: objective of the parameters; returns a real scalar.
  • gradient: function of the parameters that returns the gradient, an array for one parameter or a tuple with one gradient per parameter. Complex gradients follow dL = Re(vdot(g, dx)).
  • *parameters: as in check_pullback.
  • directions: as in check_pullback.
  • step: as in check_pullback.
  • rtol: as in check_pullback.
  • atol: as in check_pullback.
  • seed: as in check_pullback.

Returns

None when every parameter passes.

Raises

  • AssertionError: as in check_pullback.
  • ValueError: as in check_pullback, and for a complex or non-scalar objective.

check_pullback

check_pullback(record: Record, *parameters: ArrayLike, directions: tuple[ArrayLike, ...] | None=None, cotangents: ArrayLike | tuple[ArrayLike, ...] | None=None, step: float=1e-06, rtol: float=1e-05, atol: float=1e-07, seed: int=0) -> None

Check the pullback of a record against central finite differences.

The check records once and calls the pullback once with the cotangents. For each parameter it compares Re(vdot(gradient, direction)) with the central difference of Re(vdot(cotangents, value)) along the direction; the other parameters stay fixed. The shifted forward calls never use their contexts.

Args

  • record: function of the dynamic inputs that returns (value, context). The context has pullback(*cotangents) or is the pullback itself; the pullback returns one gradient per parameter, a tuple for several. The value is an array, a scalar or a tuple of them. Bind labels and options with a closure.
  • *parameters: the dynamic inputs: nonempty, finite, float64 or complex128.
  • directions: tuple with one direction per parameter, each nonzero. Default: random directions drawn from seed.
  • cotangents: gradient of the loss with respect to the value, one array or a tuple like the value. Default: random arrays drawn from seed.
  • step: finite-difference step; each parameter moves by +/- step * direction. Keep both points inside the physical domain and scale the step to the problem.
  • rtol: relative tolerance on the finite difference.
  • atol: absolute tolerance; a parameter passes when |pullback - finite difference| <= atol + rtol * |finite difference|.
  • seed: seed of the random directions and cotangents. Repeat with other seeds for more coverage.

Returns

None when every parameter passes.

Raises

  • AssertionError: a pullback disagrees; the message names the parameter, both values and the error.
  • ValueError: the record or the arguments break these rules, for example a wrong number of gradients, a wrong shape or non-finite values.

check_pushforward

check_pushforward(record: Record, *parameters: ArrayLike, directions: tuple[ArrayLike, ...] | None=None, cotangents: ArrayLike | tuple[ArrayLike, ...] | None=None, step: float=1e-06, rtol: float=1e-05, atol: float=1e-07, adjoint_rtol: float=1e-12, adjoint_atol: float=1e-12, seed: int=0) -> None

Check a record's analytic JVP against differences and its adjoint.

All parameters move together along directions. Each complete output tangent is compared with a central difference, then the identity Re(vdot(cotangent, JVP)) = Re(vdot(VJP, direction)) is checked with a fresh context. Four forward calls are made, with one derivative action on each context. No Jacobian is constructed.

Args

  • record: function of the dynamic inputs returning (value, context). The value is one array or a flat tuple of arrays; the context has pushforward(*directions) and pullback(*cotangents) methods. Its pushforward has the same output structure as the value.
  • *parameters: nonempty, finite float64 or complex128 dynamic inputs.
  • directions: tuple with one direction per parameter. Individual directions may be zero; at least one must be nonzero. Default: random directions drawn from seed.
  • cotangents: probes for the adjoint identity, with the same array or tuple structure as the output. Default: random probes from seed.
  • step: positive finite difference step. Both shifted points must remain in the physical domain.
  • rtol: relative tolerance for each finite-difference output entry.
  • atol: absolute tolerance for each finite-difference output entry.
  • adjoint_rtol: relative tolerance for the analytic adjoint identity.
  • adjoint_atol: absolute tolerance for the analytic adjoint identity.
  • seed: random seed, as in check_pullback.

Returns

None when both derivative checks pass.

Raises

  • AssertionError: a tangent or the adjoint pairing disagrees.
  • ValueError: an input, probe or derivative breaks the shape, structure or finiteness contract.
  • NotImplementedError: the record context has no pushforward method.