jax4j / NablaTensor

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

final class

Transforms

Composable tensor autodiff. Rules execute on the input backend, including enclosing tracers.

View source ↗Full module reference on javadoc.io ↗

Methods

static Tensor grad(UnaryOperator<Tensor> body, Tensor input)
static UnaryOperator<Tensor> grad(UnaryOperator<Tensor> body)
static ValueAndGrad valueAndGrad(UnaryOperator<Tensor> body, Tensor input)
static Jvp jvp(UnaryOperator<Tensor> body, Tensor input, Tensor tangent)
static ValueAndGradTree<P> valueAndGradTree(Function<P, Tensor> body, P input)
static P gradTree(Function<P, Tensor> body, P input)
static JvpTree<R> jvpTree(Function<P, R> body, P inputs, P tangents)

Forward-mode over matching tensor pytrees.

static Vjp<R, P> vjp(Function<P, R> body, P inputs, R cotangents)

Structured VJP evaluated at explicit output cotangents.

static Tensor jacfwd(UnaryOperator<Tensor> body, Tensor input)

Forward Jacobian, output dimensions followed by input dimensions. Basis sweeps are explicit.

static Tensor jacrev(UnaryOperator<Tensor> body, Tensor input)

Reverse Jacobian; one pullback per output coordinate.

static Tensor hessian(UnaryOperator<Tensor> body, Tensor input)