jax4j / NablaTensor

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

final class

CustomDerivative

A pure tensor primitive with an explicit derivative contract. Callbacks remain transformable.

View source ↗Full module reference on javadoc.io ↗

Methods

static CustomDerivative jvp(UnaryOperator<Tensor> forward, BiFunction<Tensor, Tensor, Tensor> rule)
static CustomDerivative vjp(UnaryOperator<Tensor> forward, BiFunction<Tensor, Tensor, Tensor> rule)
static CustomDerivative checkpoint(UnaryOperator<Tensor> forward)

Recompute the pure forward function during the pullback instead of retaining its internal tape.

Tensor apply(Tensor input)
DeviceBuffer evaluate(ComputeBackend backend, DeviceBuffer input)

Backend entry point; applications use apply.