Docs / jax4j-core / com.nablatensor.jax4j.autodiff

final class

AutoDiff

Reverse-mode automatic differentiation over flat or structured parameters, mirroring how the Python notebooks call jax.value_and_grad(loss)(params, x, y). Other arguments (batch data, constants) are simply closed over by lossFn.

Every intermediate Var built while lossFn runs is closed automatically once gradients have been read out, so callers only need to manage the lifetime of the returned value/gradient tensors themselves.

View source ↗Full module reference on javadoc.io ↗

Methods

static ValueAndGradTree<P> valueAndGrad(Function<? super V, Var> lossFn, P paramTree)

Computes a value and gradient for a structured parameter tree. Generic records allow the parameter and variable trees to share a runtime shape while using different leaf types, for example Net<Tensor> and Net<Var>.

static ValueAndGradTree<P> valueAndGrad(Function<? super V, Var> lossFn, P paramTree, Function<? super List<Var>, ? extends V> variableTreeFactory)

Computes a value and gradient for a structured parameter tree.

The factory receives parameter variables in deterministic pytree order and returns the tree passed to lossFn. This keeps generic records type-safe: a Net<Tensor> can produce a Net<Var> without an unchecked cast. The returned gradients have the original parameter tree's runtime shape and Tensor leaves.

A record containing parameter leaves must use a generic leaf component (for example, record Net<T>(T weight, T bias)) so the factory can create the corresponding Net<Var>.

static ValueAndGrad valueAndGrad(Function<List<Var>, Var> lossFn, List<Tensor> params)

static List<Tensor> grad(Function<List<Var>, Var> lossFn, List<Tensor> params)

Gradient-only convenience over valueAndGrad; discards the loss value.