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

final class

Jit

Traces a fluent Tensor function and fuses each maximal elementwise chain (add/sub/mul/div/unary/scalar ops) between matmul/transpose boundaries into a single backend kernel launch, instead of one launch per op. Tracing itself is cheap (plain object graph building); backends cache compiled fused kernels by their generated source, so repeated calls with the same op chain and shape reuse the compiled kernel.

View source ↗Full module reference on javadoc.io ↗

Methods

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

static Tensor apply(java.util.function.Function<java.util.List<Tensor>, Tensor> body, java.util.List<Tensor> inputs)

Captures all arguments so layers, parameters and optimizer state remain fusible.

static Compiled compile(UnaryOperator<Tensor> body)

Captures a pure function once per shape/dtype/device/backend signature. Runtime values remain inputs. Closed-over tensors must remain alive; Java side effects and scalar constants are evaluated only during capture.

static CompiledTree<P, R> compileTree(java.util.function.Function<P, R> body)

Captures a Tensor-leaf pytree and preserves the output tree, including value/gradient records.