Module Rune

Parameter trees: Nx.Ptree.S is the traversal interface every transformation takes, and Nx.Ptree.t the stock dynamic instance.

grad (module P) f params is the gradient of f at params, with the same structure and leaf types as params. Leaves of params that do not contribute to the result have all-zero gradients.

Raises Invalid_argument if a parameter leaf has an integer or boolean dtype (hold non-differentiable data in the closure or the auxiliary output), or if f params is not a scalar (a tensor with exactly one element); use vjp to differentiate non-scalar outputs against an explicit cotangent.

value_and_grad (module P) f params is (f params, grad (module P) f params), computed in a single forward and backward pass.

value_and_grad_aux (module P) f params is like value_and_grad for an objective returning auxiliary data alongside its result. The auxiliary value is returned as-is and does not contribute to the gradient.

vjp (module P) f params cotangent is (f params, grads) where grads is the vector-Jacobian product of f at params against cotangent, with the same structure and leaf types as params. cotangent must have f params's shape and dtype.

vjp2 (module P) (module Q) f params cotangents is like vjp for an objective returning a structured output: cotangents provides one cotangent per output leaf, each with its output leaf's shape and dtype, and the pulled-back cotangents have params' structure.

Raises Invalid_argument if a cotangent leaf's shape does not match its output leaf's shape.

vjp_fun (module P) f params is (f params, pullback). pullback ct is the vector-Jacobian product of f at params against ct; it may be called any number of times with different cotangents, each call running one backward pass over the recorded computation without re-running f. Calling the pullback under another transformation (for example vmap) transforms the backward pass. Pullbacks are not thread-safe.

Raises Invalid_argument if ct's shape does not match the output's.

vjp_fun' f x is like vjp_fun for a function of a single tensor.

jvp (module P) f params tangents is (f params, df) where df is the Jacobian-vector product of f at params against tangents, computed in a single forward pass. tangents must be structurally equal to params; each tangent leaf must have its parameter leaf's shape. The output may have any shape.

Raises Invalid_argument if a tangent leaf's shape does not match its parameter leaf's shape.

jvp_aux (module P) f params tangents is like jvp for an objective returning auxiliary data alongside its result. The auxiliary value is returned as-is and does not contribute to the tangent.

jvp2 (module P) (module Q) f params tangents is like jvp for an objective returning a structured output: the result tangent has the output's structure, one tangent per output leaf.

vmap ?in_axes ?out_axis (module P) f params maps f over the tensor leaves of params. f is written for unbatched values: it observes each mapped leaf without its mapped axis, and its result gains a batch axis at out_axis (default 0). Values f closes over are constants of the map, and a result that does not depend on the mapped inputs is broadcast along the batch axis.

in_axes gives the mapped axis of each leaf, paired with leaves in the structure's traversal order: Some i maps axis i (negative from the end), None passes the leaf whole as a constant. It defaults to mapping axis 0 of every leaf. Mapped axes must agree on their size.

Composes with the other transformations: vmap of grad computes per-example gradients, and grad of vmap differentiates through the map.

Note. Randomness a lane closes over (Nx.rand and friends, or a Rng key used as a constant) draws identical values for every lane: it is a constant of the map. For decorrelated lanes, Rng.split one key into per-lane keys, stack them into an [n; 2] tensor, and map over it — the sampler inside the map then sees one key per lane. Reading a batched tensor's value inside the mapped function raises.

Raises Invalid_argument if in_axes does not have one entry per leaf, maps no leaf, names an axis out of bounds, or if the mapped axis sizes disagree.

vmap2 ?in_axes ?out_axis (module P) (module Q) f params is like vmap for a mapped function returning a structured output: every output leaf gains a batch axis at out_axis, and output leaves that do not depend on the mapped inputs are broadcast.

vmap' ?in_axis ?out_axis f x is like vmap for a function of a single tensor, mapping over x's axis in_axis and placing the batch axis of the result at out_axis. Both default to 0.

custom_vjp (module P) ~fwd ~bwd params is fst (fwd params), with a user-defined reverse rule. Under the innermost reverse-mode transformation, fwd's internal operations are not differentiated; instead bwd residual cotangent provides the parameter gradients, with each leaf matching its parameter leaf's shape and dtype. residual is whatever fwd returned alongside its result. Enclosing transformations (an outer grad, vmap) see the forward computation itself.

Raises Invalid_argument if the call is differentiated in forward mode; define a custom_jvp rule for that.

custom_jvp (module P) ~f ~jvp params is f params, with a user-defined forward rule. Under the innermost forward-mode transformation, jvp params tangents provides both the result and its tangent, replacing f's internal operations.

Raises Invalid_argument if the call is differentiated in reverse mode; define a custom_vjp rule for that.

grad' f x is like grad for a function of a single tensor.

value_and_grad' f x is like value_and_grad for a function of a single tensor.

vjp' f x cotangent is like vjp for a function of a single tensor.

jvp' f x tangent is like jvp for a function of a single tensor.

remat (module P) f params is f params, recomputed during the backward pass instead of having its intermediate results retained by the tape: reverse-mode differentiation of a remated function trades compute for memory. Gradients are unchanged.

Raises Invalid_argument if differentiated in forward mode (it is a custom_vjp rule underneath).

jacfwd' f x is the Jacobian of f at x, with shape shape (f x) @ shape x, computed column by column in forward mode (one vectorized pass). Prefer it when the input is smaller than the output.

jacrev' f x is the Jacobian of f at x, with shape shape (f x) @ shape x, computed row by row in reverse mode (one forward pass, one vectorized backward pass). Prefer it when the output is smaller than the input.

hessian' f x is the Hessian of the scalar objective f at x, with shape shape x @ shape x (forward over reverse).

hvp (module P) f params v is the Hessian-vector product of the scalar objective f at params against v, with params' structure, computed without materializing the Hessian (forward over reverse).

hvp' f x v is like hvp for a function of a single tensor.

check_grads (module P) f params compares the reverse-mode gradient of the scalar objective f at params against central-difference directional derivatives along deterministic directions. Ok () means they agree within tol (relative, default 1e-2); Error msg describes the disagreement. eps is the finite-difference step (default 1e-4).

The check is directional, not per-element: it validates gradients cheaply rather than exhaustively. Use float64 parameters for reliable results; float32 may need a looser tol.

Explicit splittable PRNG keys.

Raised when a function cannot be compiled: it read the value of a traced tensor (for example Nx.item on a value that depends on the inputs, or a data-dependent branch), it assigned to a tensor it closes over (captures are compile-time constants), it drew random values from a key that does not depend on the inputs (Nx.rand and friends, or a captured Rng.key — the draw would be a compile-time constant replayed on every call; pass the key as an input instead), or it used an operation the compiler does not support (FFT, linear algebra, complex, int4 and uint4 tensors, assigning into a view).

jit (module P) f is f compiled. The first application traces f, compiles the traced computation into fused kernels, and runs them; later applications with the same leaf signature — dtypes and shapes, in traversal order — replay the compiled program on the new leaf values. A new signature triggers a fresh trace and compilation.

device selects where the kernels compile and run: "CPU" (the default), "CUDA" (NVIDIA GPUs), or "METAL" (macOS only). On the CPU device, contiguous inputs and captured tensors are read in place and outputs are computed directly into the returned tensors' storage; non-contiguous tensors are copied.

On other devices, results are bit-identical but data moves lazily. Inputs are copied to the device on every call; outputs stay resident on the device and materialize on the host the first time their data is read (metadata such as shape and dtype never transfers). An unread output fed back as an input leaf of any jit call on the same device seeds the compiled program's input directly — no transfer — which makes iterated calls (training steps, decode loops with a cache) run without per-call traffic. Reading an output moves it to the host for good: it becomes an ordinary tensor, later mutations are honored, and feeding it back uploads its current bytes. Device memory backing an output is held until the output is read or garbage-collected; past a budget (the RUNE_JIT_RESIDENT_BUDGET environment variable, in bytes, a few GiB by default) a collection is triggered before allocating more. A transfer failure surfaces as an exception at the first read of the affected output.

donate (default false) consumes the resident inputs: once a call completes — never during it — the device buffers of every input leaf that was an unread resident output of an earlier call are released to the allocator, and the donated handle becomes unusable — reading it, or feeding it to a later call (which reads it), raises Invalid_argument; read or copy the value before the call if it is still needed. The next call's fresh outputs reuse the released storage, so a state-to-state loop (state <- step state with step jitted ~donate:true) holds about two generations of state on the device instead of one per call awaiting collection. Only resident handles are consumed: host tensors and handles already read are unaffected. A handle appearing as several input leaves is donated once, and a leaf the function updates in place is materialized on the host by the writeback rather than donated.

The compilation cache lives in the partial application jit (module P) f: apply jit once and reuse the returned function. Tensors f closes over are compile-time constants, bound once when the trace first compiles: on the CPU device contiguous captures are read in place, and every other capture is copied to the device once per closure — signatures share the copy. Mutating a captured tensor between calls is not supported and has unspecified visibility (the CPU device may observe the mutation through its in-place binding; other devices never do): pass values that change between calls as leaves of P rather than capturing them.

Compiled programs also persist across processes: the first compilation of a trace writes the scheduled and compiled kernels to a disk cache under the platform cache directory ($XDG_CACHE_HOME/tolk/rune_jit, with XDG_CACHE_HOME defaulting to ~/.cache on Linux and ~/Library/Caches on macOS), and a later process compiling the same trace loads them instead of scheduling and compiling again, leaving tracing as the bulk of warm start-up time. Entries are invalidated automatically when the executable, the device or its compiler, or the codegen options change. Set the JITCACHE environment variable to 0 to disable the persistent cache; pmap compilations are never persisted. Results are identical either way.

Under an enclosing transformation (grad, vmap, with_debug, an outer jit), the wrapped function runs directly so the transformation observes its operations: jit never changes results, only speed. Compose the other way — differentiate inside the jitted function — to compile the forward and backward passes together:

let train_step =
  Rune.jit2
    (module Params)
    (module Params)
    (fun p ->
      let g = Rune.grad (module Params) loss p in
      Params.map2 (fun w g -> Nx.sub w (Nx.mul_s g lr)) p g)

Whole-tensor in-place updates (Nx.assign or Nx.blit on an input leaf) are replayed by writing the computed value back into the destination, so in-place state threaded through P carries across calls on every device. Assigning to a captured tensor raises Jit_error at trace time — captures are compile-time constants; thread mutable state through the input structure instead. Structured values read during tracing must not depend on traced tensors: a data-dependent cond or while_loop predicate raises Jit_error. Compiled functions are not thread-safe.

Randomness inside a jitted function comes from Rng keys threaded through the inputs: samplers are pure functions of their key, so the compiled program recomputes each draw from the current key on every call — feed a fresh key (Rng.split, Rng.fold_in) for fresh values. A key that does not depend on the inputs (implicit RNG such as Nx.rand, or a captured key) raises Jit_error at trace time rather than replay one frozen draw.

Raises Jit_error when tracing fails (Jit_error), and Invalid_argument for an unknown or unavailable device.

jit2 (module P) (module Q) f is like jit for a function returning a structured output.

jit' f is like jit for a function of a single tensor.

pmap ~devices (module P) f is f compiled to run in parallel across devices — jit whose inputs are placed on a device tuple instead of one device. Device names are as in jit, with an instance suffix to address several devices of one backend ("CUDA:0", "CUDA:1", or "CPU:1", "CPU:2", ...); all devices must share one backend.

in_axes gives one entry per leaf of P, in traversal order: Some a splits the leaf along axis a into List.length devices equal shards, one per device; None replicates the leaf, a full copy on every device. It defaults to Some 0 for every leaf. f observes full (global) shapes and needs no collective operations: an operation combining sharded and replicated values runs on every device over its shard, and a reduction over a sharded axis (say the mean loss over a sharded batch) becomes a cross-device allreduce automatically — differentiating such a loss inside pmap yields allreduced gradients, which makes data-parallel training a matter of sharding the batch and replicating the parameters.

Compilation, caching, capture, and writeback semantics are jit's. Outputs stay resident, one buffer per device, and gather to the host the first time they are read (shards are reassembled along their axis; replicated outputs read one replica). An unread output fed back as an input leaf whose placement matches — same devices, same axis or replication — seeds the compiled program's buffers directly with no transfer, so iterated calls (a data-parallel training step) move only the freshly sharded batch. donate consumes resident inputs as in jit, releasing every per-device buffer of the donated handle; a handle whose placement mismatches is forced to the host first and is not donated. Assigning a sharded value to an input leaf raises Jit_error at trace time (it would gather on every call); replicated writebacks are honored.

Under an enclosing transformation, f runs directly on the host like jit: differentiate inside the pmapped function.

Raises Invalid_argument if devices is empty, mixes backends, or a device is unavailable; if in_axes has one entry per leaf missing or in excess; or if a sharded leaf's dimension does not divide evenly across the devices. Raises Jit_error when tracing fails, as jit.

pmap2 (module P) (module Q) f is like pmap for a function returning a structured output.

Transfer accounting for compiled functions. The zero-copy CPU path moves no bytes and counts nothing.

jit_stats () is the current transfer counters, cumulative over the whole program. Set the RUNE_JIT_DEBUG environment variable to 1 to also log a per-call summary to stderr.

reset_jit_stats () zeroes the cumulative transfer counters. resident_bytes tracks live state and is not reset.

scan (module C) ~f ~init xs folds f over slices of xs along axis 0: f carry x returns the next carry and a per-step output. The result is the final carry and the outputs stacked along a new axis 0. Differentiating traces every step.

Raises Invalid_argument if xs is a scalar or empty along axis 0.

cond pred ~then_ ~else_ runs one branch according to the scalar pred. Reading pred concretizes it: inside vmap, a predicate that depends on the mapped inputs raises, since the lanes could diverge.

while_loop (module C) ~cond ~body init iterates body on the carry while cond holds. Reading the predicate concretizes it, with the same vmap caveat as cond. Differentiating traces every iteration actually taken.

with_debug f runs f and logs each tensor operation it performs — the operation name and output shape — to ppf (defaults to Format.err_formatter). Composes with the other transformations: run it outermost to also observe the operations they emit. Uncommon operations may execute unlogged.

detach t is a copy of t through which gradients do not flow. Use it to hold a value constant inside a differentiated function, including as input to an operation whose gradient is not implemented.

no_grad f runs f with gradient tracking disabled: tensors it produces are constants of the surrounding differentiation.