Struct Func 

pub struct Func { /* private fields */ }

A function of named inputs and parameters, traced into a graph. Unit of composition for the grad/vmap/ jvp/hvp transforms.

Inputs vs parameters. Declare per-call data with GraphScope::input and trainable weights with GraphScope::param. Bind weight values once with with_param; they are applied on every run and survive transforms — so f.grad(&["w"]) gives gradients w.r.t. the weights (training), runnable with just the inputs.

Implementations

impl Func

pub fn new( name: impl Into<String>, build: impl FnOnce(&mut GraphScope) -> Tensor, ) -> Self

Trace a single-output function. The closure builds the body from inputs declared via GraphScope::input; the returned tensor is the output.

pub fn multi( name: impl Into<String>, build: impl FnOnce(&mut GraphScope) -> Vec<Tensor>, ) -> Self

Trace a multi-output function (the closure returns all outputs in order).

pub fn from_graph(graph: Graph) -> Self

Wrap an already-built graph (its outputs are this function’s outputs).

pub fn with_param( self, name: impl Into<String>, data: impl Into<Vec<f32>>, ) -> Self

Bind a parameter (weight) value by name. Overwrites any prior binding for that name. Builder-style: chain several, then run/grad/jit.

pub fn graph(&self) -> &Graph

The underlying graph.

pub fn output_count(&self) -> usize

Number of outputs this function produces.

pub fn param_binding(&self, name: &str) -> Option<&[f32]>

The currently bound value of a parameter, if any. After train_step this is the updated weight.

pub fn param_names(&self) -> Vec<String>

The names of every trainable parameter (Op::Param) in the graph, in declaration order — exactly the wrt list a full-model training step needs, so you never hand-maintain it. Powers init_params, the *_all training steps, and checkpoint round-trips.

pub fn param_shape_of(&self, name: &str) -> Option<Vec<usize>>

The static shape (dims) of a parameter by name, or None if there is no such param. A dynamic axis reads as 0.

pub fn init_params(self, init: impl FnMut(&str, &[usize]) -> Vec<f32>) -> Self

Bind every declared parameter by calling init(name, &static_dims) and using the returned Vec<f32> — one call seeds a whole model’s weights, instead of a with_param per tensor. Overwrites any existing bindings. Builder-style: chain before train_step.

let model = Func::from_graph(rlx! { … }).init_params(|name, dims| {
    if name.ends_with(".bias") { vec![0.0; dims.iter().product()] }
    else { he_init(dims) }
});

pub fn init_randn(self, seed: u64, stddev: f32) -> Self

Seed every parameter with i.i.d. Gaussian noise N(0, stddev²) from a deterministic SplitMix64 + Box–Muller stream keyed by seed (reproducible, no rand dependency). A one-liner for a from-scratch model; for per-tensor schemes (zeroed biases, fan-in scaling) pass your own closure to init_params.

impl Func

Composable transforms (require rlx_autodiff; enabled by grad / transforms / autodiff).

pub fn grad(&self, wrt: &[&str]) -> Func

Reverse-mode gradient w.r.t. the named inputs. The first output is treated as the (scalar) loss; the unit seed is baked in, so the result is a function inputs -> [∂loss/∂wrt …] runnable with no extra feeds.

pub fn value_and_grad(&self, wrt: &[&str]) -> Func

Like grad but the result emits the loss and the gradients: outputs [loss, ∂loss/∂wrt[0], …]. The JAX value_and_grad.

pub fn value_and_grad_all(&self) -> Func

value_and_grad w.r.t. every parameter in the graph — no hand-maintained name list. Outputs [loss, ∂loss/∂p …] ordered like param_names.

pub fn grad_all(&self) -> Func

grad w.r.t. every parameter in the graph.

pub fn vmap(&self, batched: &[&str], batch_size: usize) -> Func

Vectorize over a leading batch axis on the named inputs (out_axes = 0). Inputs not listed are shared across the batch.

pub fn jvp(&self, tangent_for: &[&str]) -> Func

Forward-mode Jacobian-vector product w.r.t. the named inputs/params. The returned function gains a tangent input per tangent_for entry.

pub fn hvp(&self, wrt: &[&str]) -> Func

Hessian-vector product w.r.t. the named inputs/params.

impl Func

Execution (requires the eval feature).

pub fn run(&self, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>>

Compile + run with named inputs, one Vec<f32> per output. The device is auto-selected (fastest compiled-in backend that can run the graph).

pub fn run_on(&self, device: Device, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>>

Compile + run on an explicit device. The compiled graph is cached, so repeated run calls with different inputs recompile nothing — a real jit. Bound params are applied before each run.

pub fn jit(&self) -> Jitted

Compile now and return a callable that holds the compiled graph, so each Jitted::run skips even the cache lookup. The JAX-shaped jit(f); composes after the other transforms: f.grad(&["x"]).vmap(&["x"], 4).jit().

pub fn jit_on(&self, device: Device) -> Jitted

Func::jit targeting an explicit device.

impl Func

Training (requires the optim feature: autodiff + eval + rlx_optim).

pub fn train_step_at( &self, opt: &mut dyn Optimizer, schedule: &LrSchedule, step: usize, wrt: &[&str], inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

One optimization step. Computes value + gradients over wrt (which must name bound params), applies opt to each param, and returns the updated function plus the loss. Drive a loop by reassigning the returned Func:

let mut opt = Sgd::new(0.1);
for _ in 0..100 {
    let (next, loss) = model.train_step(&mut opt, &["w"], &[("x", &xs)]);
    model = next;
}

train_step with a learning-rate schedule: sets opt’s LR to schedule.lr_at(step) before stepping. Drive a loop by passing the iteration index as step.

pub fn train_step( &self, opt: &mut dyn Optimizer, wrt: &[&str], inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

pub fn train_step_on( &self, device: Device, opt: &mut dyn Optimizer, wrt: &[&str], inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

train_step pinned to an explicit device instead of the auto-selected fastest backend — the one call you need to train a whole run on, say, Metal (no hand-rolled value_and_grad().run_on() + optimizer loop).

pub fn train_step_at_on( &self, device: Device, opt: &mut dyn Optimizer, schedule: &LrSchedule, step: usize, wrt: &[&str], inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

train_step_on with a learning-rate schedule.

pub fn train_step_all( &self, opt: &mut dyn Optimizer, inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

train_step w.r.t. every parameter — the wrt list is param_names, so the call is just (opt, inputs).

pub fn train_step_all_at( &self, opt: &mut dyn Optimizer, schedule: &LrSchedule, step: usize, inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

All-parameter train_step_at.

pub fn train_step_all_on( &self, device: Device, opt: &mut dyn Optimizer, inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

All-parameter train_step_on.

pub fn train_step_all_at_on( &self, device: Device, opt: &mut dyn Optimizer, schedule: &LrSchedule, step: usize, inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

The full-model, device-pinned, scheduled training step — the whole loop body for a from-scratch model is model.train_step_all_at_on(dev, &mut opt, &sched, step, feed).

pub fn train_step_all_at_on_clipped( &self, device: Device, opt: &mut dyn Optimizer, schedule: &LrSchedule, step: usize, max_grad_norm: f32, inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

train_step_all_at_on with global-L2-norm gradient clipping to max_grad_norm — the standard cure for the loss spikes that can blow a run up to NaN late in training (~1.0 is typical). A non-positive max_grad_norm disables clipping (identical to the unclipped variant).

pub fn train_step_all_at_on_qat( &self, device: Device, opt: &mut dyn Optimizer, schedule: &LrSchedule, step: usize, max_grad_norm: f32, quant: impl Fn(&mut [f32]), inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

Quantization-aware full-model step: each parameter is passed through quant (an in-place quantizer) before the forward, and the gradient is taken at that quantized point — but applied to the f32 master weights (a straight-through estimator). This trains the model at whatever emulated precision quant imposes, on any backend. Pair with crate::lowp for fXmYeZ / nvf4 / f8 / bf8 formats:

let (e, m, max) = rlx_tensor::lowp::parse_format("nvf4").unwrap();
let (next, loss) = model.train_step_all_at_on_qat(
    dev, &mut opt, &sched, step, 1.0,
    |w| rlx_tensor::lowp::quantize_slice(w, e, m, max), feed);

pub fn train_step_all_at_on_qat_with( &self, device: Device, opt: &mut dyn Optimizer, schedule: &LrSchedule, step: usize, max_grad_norm: f32, quant: impl Fn(&str, &[usize], &mut [f32]), inputs: &[(&str, &[f32])], ) -> (Func, Vec<f32>)

train_step_all_at_on_qat where the quantizer also sees each parameter’s name and shape.

Block formats need this. A per-block scale is only usable by hardware if its blocks run along the axis a dot product contracts over, and a flat &mut [f32] does not say which axis that is — blocking the wrong way gives every output its own scale mid-dot-product, which no kernel can apply. Per-tensor and per-element formats can keep using the simpler entry point.

Trait Implementations

impl Clone for Func

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Func

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.