Struct Graph 

pub struct Graph {
    pub name: String,
    pub outputs: Vec<NodeId>,
    /* private fields */
}

A computation graph — the core IR data structure.

Example

use rlx_ir::*;

let mut g = Graph::new("bert_layer");

// Inputs
let x = g.input("hidden", Shape::new(&[4, 15, 384], DType::F32));
let w = g.param("qkv_weight", Shape::new(&[384, 1152], DType::F32));
let b = g.param("qkv_bias", Shape::new(&[1152], DType::F32));

// QKV projection: matmul + bias
let mm = g.matmul(x, w, Shape::new(&[4, 15, 1152], DType::F32));
let qkv = g.binary(op::BinaryOp::Add, mm, b, Shape::new(&[4, 15, 1152], DType::F32));

assert_eq!(g.len(), 5);
println!("{g}");

Fields

name: String

outputs: Vec<NodeId>

Output node IDs (the graph’s results).

Implementations

impl Graph

pub fn new(name: impl Into<String>) -> Graph

pub fn len(&self) -> usize

Number of nodes in the graph.

pub fn is_empty(&self) -> bool

pub fn node(&self, id: NodeId) -> &Node

Get a node by ID.

pub fn nodes(&self) -> &[Node]

Iterate all nodes in topological order (insertion order = topo order).

pub fn shape(&self, id: NodeId) -> &Shape

Get the shape of a node’s output.

pub fn set_outputs(&mut self, outputs: Vec<NodeId>)

Set the graph outputs.

pub fn set_inputs(&mut self, id: NodeId, inputs: Vec<NodeId>)

Replace the input list of a node in place. Used by post- construction passes (quant_propagate, dce, etc.) that rewire consumers without inserting new nodes. Caller is responsible for shape consistency — this does no re-inference.

pub fn node_mut(&mut self, id: NodeId) -> &mut Node

pub fn nodes_mut(&mut self) -> &mut [Node]

pub fn append_node( &mut self, op: Op, inputs: Vec<NodeId>, shape: Shape, name: Option<String>, ) -> NodeId

Append a node to the graph. pub(crate) so per-op builder files in rlx_ir::ops::* can call it (plan #53). Append a node for backend graph slicing (e.g. TPU HLO segments).

pub fn users(&self, id: NodeId) -> Vec<NodeId>

Find all nodes that use a given node’s output.

pub fn use_count(&self, id: NodeId) -> usize

Count how many nodes use a given node’s output.

pub fn node_id_by_name(&self, name: &str) -> Option<NodeId>

Find a node by the name on its Op::Input or Op::Param.

Input/Param leaves are the graph’s roots and survive optimizer passes (fusion, DCE), so a name is the stable way to recover a handle into a rewritten graph whose NodeIds have been renumbered — e.g. after fuse. Outputs are already positionally stable in outputs. Returns the first match (names are expected unique).

pub fn input_id(&self, name: &str) -> Option<NodeId>

node_id_by_name restricted to graph inputs.

pub fn param_id(&self, name: &str) -> Option<NodeId>

node_id_by_name restricted to parameters.

pub fn topo_order(&self) -> impl Iterator<Item = NodeId>

Topological order (already guaranteed by construction — just node indices).

pub fn reverse_topo(&self) -> impl Iterator<Item = NodeId>

Reverse topological order (outputs first).

pub fn define( name: impl Into<String>, build: impl FnOnce(&mut HirModule) -> HirNodeId, ) -> GraphModule

Fusion-first model definition at HIR level.

Returns a [GraphModule] at HIR stage; call [GraphModule::lower] or pass to [rlx_opt::CompilePipeline::compile_module].

pub fn hir(name: impl Into<String>) -> GraphModule

Start an empty HIR-stage [GraphModule].

pub fn module(self) -> GraphModule

Wrap this MIR graph in a [GraphModule] for pipeline operations.

pub fn from_hir(hir: HirModule) -> Result<Graph, LowerError>

Lower a HIR module to a MIR graph.

pub fn to_mir(self) -> MirModule

View as [MirModule].

pub fn from_lir(lir: LirModule) -> Graph

Extract the MIR graph from optimized LIR.

pub fn inspect(&self) -> String

Annotated text dump ([inspect_graph]).

pub fn has_dynamic_dims(&self) -> bool

True if any node shape uses a [Dim::Dynamic] symbol.

pub fn dynamic_symbols(&self) -> Vec<u32>

All dynamic symbols referenced in this graph.

pub fn bind(&self, bindings: &DimBinding) -> Graph

Specialize symbolic dims to concrete sizes.

pub fn inspect_module(module: &GraphModule) -> String

Stage-aware dump when wrapped in [GraphModule].

impl Graph

pub fn attention( &mut self, q: NodeId, k: NodeId, v: NodeId, mask: NodeId, num_heads: usize, head_dim: usize, shape: Shape, ) -> NodeId

Scaled dot-product attention with a custom (caller-supplied) mask. Equivalent to attention_kind(.., MaskKind::Custom, ..).

pub fn attention_opts( &mut self, q: NodeId, k: NodeId, v: NodeId, mask: NodeId, num_heads: usize, head_dim: usize, shape: Shape, score_scale: Option<f32>, attn_logit_softcap: Option<f32>, ) -> NodeId

Like Self::attention with optional score scale and logit softcap.

pub fn attention_kind( &mut self, q: NodeId, k: NodeId, v: NodeId, num_heads: usize, head_dim: usize, mask_kind: MaskKind, shape: Shape, ) -> NodeId

Scaled dot-product attention with a kernel-synthesized mask (None / Causal / SlidingWindow). Inputs are Q, K, V only — no mask tensor is allocated or read in the inner loop. Use MaskKind::None for a single un-padded sequence.

pub fn attention_kind_opts( &mut self, q: NodeId, k: NodeId, v: NodeId, num_heads: usize, head_dim: usize, mask_kind: MaskKind, shape: Shape, score_scale: Option<f32>, attn_logit_softcap: Option<f32>, ) -> NodeId

Like Self::attention_kind with optional score scale and logit softcap.

pub fn attention_bias( &mut self, q: NodeId, k: NodeId, v: NodeId, bias: NodeId, num_heads: usize, head_dim: usize, shape: Shape, ) -> NodeId

Scaled dot-product attention with an additive bias tensor of shape [batch, num_heads, query_len, key_len] added to the QK^T · scale scores before softmax. Lets boxRPB / per-query position biases reuse the fast Op::Attention kernel path.

pub fn attention_bias_opts( &mut self, q: NodeId, k: NodeId, v: NodeId, bias: NodeId, num_heads: usize, head_dim: usize, shape: Shape, score_scale: Option<f32>, attn_logit_softcap: Option<f32>, ) -> NodeId

Like Self::attention_bias with optional score scale and logit softcap (e.g. Gemma’s score_scale = 1.0 + final-attn-logit softcap, which the bias-less variant would silently drop).

impl Graph

pub fn log_mel(&mut self, spectrum: NodeId, filters: NodeId) -> NodeId

Log-mel spectrogram (Whisper-style) from RLX FFT block-layout spectrum.

• spectrum — [..., 2*n_fft] with re plane then im plane (same as Op::Fft output).
• filters — [n_mels, n_bins] mel filterbank (n_bins = n_fft/2 + 1).

Output: [..., n_mels].

pub fn log_mel_backward( &mut self, spectrum: NodeId, filters: NodeId, dy: NodeId, ) -> NodeId

VJP of [log_mel] w.r.t. spectrum.

pub fn welch_peaks( &mut self, spectrum: NodeId, k: usize, n_segments: usize, ) -> NodeId

Top-K Welch peaks from RLX FFT block-layout segment spectra.

• spectrum — [batch * n_segments, 2*n_fft]
• k — spikes per batch row
• n_segments — Welch segments averaged per row

Output: [batch, k*2] interleaved (bin, power).

impl Graph

pub fn axial_rope2d( &mut self, x: NodeId, end_x: usize, end_y: usize, head_dim: usize, num_heads: usize, theta: f32, repeat_factor: usize, ) -> NodeId

x: [1, seq, num_heads * head_dim] → same shape.

impl Graph

pub fn relu_backward(&mut self, x: NodeId, dy: NodeId) -> NodeId

ReLU backward: dx = dy where x > 0 else 0. Output shape matches x.

pub fn activation_backward( &mut self, kind: Activation, x: NodeId, dy: NodeId, ) -> NodeId

Element-wise activation backward — closed-form derivative of any single-input activation other than ReLU. See Op::ActivationBackward for the per-kind formulae.

pub fn layer_norm_backward_input( &mut self, x: NodeId, gamma: NodeId, dy: NodeId, axis: i32, eps: f32, ) -> NodeId

LayerNorm backward w.r.t. the input. Inputs [x, gamma, dy]. Output shape matches x. Currently axis = -1 only.

pub fn rms_norm_backward_input( &mut self, x: NodeId, gamma: NodeId, beta: NodeId, dy: NodeId, axis: i32, eps: f32, ) -> NodeId

RMSNorm backward w.r.t. input. Inputs [x, gamma, beta, dy].

pub fn rms_norm_backward_gamma( &mut self, x: NodeId, gamma: NodeId, beta: NodeId, dy: NodeId, axis: i32, eps: f32, ) -> NodeId

pub fn rms_norm_backward_beta( &mut self, x: NodeId, gamma: NodeId, beta: NodeId, dy: NodeId, axis: i32, eps: f32, ) -> NodeId

pub fn rope_backward( &mut self, dy: NodeId, cos: NodeId, sin: NodeId, head_dim: usize, n_rot: usize, ) -> NodeId

pub fn cumsum_backward( &mut self, dy: NodeId, out_shape: Shape, axis: i32, exclusive: bool, ) -> NodeId

pub fn gather_backward( &mut self, dy: NodeId, indices: NodeId, table_shape: Shape, axis: i32, ) -> NodeId

pub fn group_norm_backward_input( &mut self, x: NodeId, gamma: NodeId, beta: NodeId, dy: NodeId, num_groups: usize, eps: f32, ) -> NodeId

GroupNorm (NCHW) backward w.r.t. input. Inputs [x, gamma, beta, dy].

pub fn group_norm_backward_gamma( &mut self, x: NodeId, dy: NodeId, gamma_shape: Shape, num_groups: usize, eps: f32, ) -> NodeId

GroupNorm backward w.r.t. gamma. Inputs [x, dy].

pub fn group_norm_backward_beta( &mut self, x: NodeId, dy: NodeId, beta_shape: Shape, num_groups: usize, eps: f32, ) -> NodeId

GroupNorm backward w.r.t. beta. Inputs [x, dy].

pub fn batch_norm_inference_backward_input( &mut self, x: NodeId, gamma: NodeId, mean: NodeId, var: NodeId, dy: NodeId, eps: f32, ) -> NodeId

BatchNorm inference backward w.r.t. input.

pub fn batch_norm_inference_backward_gamma( &mut self, x: NodeId, mean: NodeId, var: NodeId, dy: NodeId, gamma_shape: Shape, eps: f32, ) -> NodeId

BatchNorm inference backward w.r.t. gamma.

pub fn batch_norm_inference_backward_beta( &mut self, dy: NodeId, beta_shape: Shape, ) -> NodeId

BatchNorm inference backward w.r.t. beta.

pub fn layer_norm_backward_gamma( &mut self, x: NodeId, dy: NodeId, gamma_shape: Shape, axis: i32, eps: f32, ) -> NodeId

LayerNorm backward w.r.t. gamma. Inputs [x, dy]. Output shape is provided by the caller — typically the gamma’s shape, e.g. [D] for a per-feature 1-D gamma.

pub fn maxpool2d_backward( &mut self, x: NodeId, dy: NodeId, kernel_size: Vec<usize>, stride: Vec<usize>, padding: Vec<usize>, ) -> NodeId

2D max-pool backward. x is the original NCHW input; dy is the upstream gradient with shape matching the pool’s output. Output shape matches x.

pub fn conv2d_backward_input( &mut self, dy: NodeId, w: NodeId, x_shape: Shape, kernel_size: Vec<usize>, stride: Vec<usize>, padding: Vec<usize>, dilation: Vec<usize>, groups: usize, ) -> NodeId

Conv2D backward w.r.t. input. dy has the conv output shape; w is the forward weight [C_out, C_in/groups, kH, kW]. The output shape (the original input shape) is supplied by the caller because it can’t be unambiguously derived from dy.shape alone in the presence of strides + padding.

pub fn conv2d_backward_weight( &mut self, x: NodeId, dy: NodeId, w_shape: Shape, kernel_size: Vec<usize>, stride: Vec<usize>, padding: Vec<usize>, dilation: Vec<usize>, groups: usize, ) -> NodeId

Conv2D backward w.r.t. weight. Output shape matches the forward weight [C_out, C_in/groups, kH, kW].

pub fn softmax_cross_entropy( &mut self, logits: NodeId, targets: NodeId, ) -> NodeId

Fused softmax + cross-entropy against a dense target distribution (soft labels / one-hot probabilities). logits [N, C], targets [N, C] → [N] per-row loss loss[n] = logsumexp(logits[n]) - Σ_c targets[n,c]·logits[n,c].

pub fn softmax_cross_entropy_with_logits( &mut self, logits: NodeId, labels: NodeId, ) -> NodeId

Fused softmax + cross-entropy with f32-encoded integer labels. logits [N, C], labels [N] → [N] per-row loss.

pub fn softmax_cross_entropy_backward( &mut self, logits: NodeId, labels: NodeId, d_loss: NodeId, ) -> NodeId

Backward of softmax_cross_entropy_with_logits. [logits, labels, d_loss] → dlogits shaped like logits.

pub fn complex_norm_sq(&mut self, z: NodeId) -> NodeId

Element-wise complex squared-magnitude: |z|² = re² + im². Input must be DType::C64; output is same logical shape but DType::F32. The canonical real-valued loss surface for Wirtinger reverse-mode AD on complex graphs.

pub fn attention_backward( &mut self, wrt: AttentionBwdWrt, q: NodeId, k: NodeId, v: NodeId, dy: NodeId, num_heads: usize, head_dim: usize, mask_kind: MaskKind, mask: Option<NodeId>, ) -> NodeId

Scaled dot-product attention backward w.r.t. q, k, or v. See Op::AttentionBackward. When mask_kind is MaskKind::Custom or MaskKind::Bias, pass the same mask tensor used in forward.

pub fn attention_backward_all( &mut self, q: NodeId, k: NodeId, v: NodeId, dy: NodeId, num_heads: usize, head_dim: usize, mask_kind: MaskKind, mask: Option<NodeId>, ) -> (NodeId, NodeId, NodeId)

Emit dQ, dK, and dV for one Op::Attention forward node.

pub fn complex_norm_sq_backward(&mut self, z: NodeId, g: NodeId) -> NodeId

Wirtinger backward for [complex_norm_sq]: given upstream g (real, same shape as the forward output) and the original complex input z, returns dz = g · z as C64.

pub fn conjugate(&mut self, z: NodeId) -> NodeId

Element-wise complex conjugate: z̄ = re - i·im. Input must be DType::C64; output is the same shape and dtype. Used by Wirtinger VJP rules on C64 binary ops.

impl Graph

pub fn linear_bias( &mut self, input: NodeId, weight: NodeId, bias: Option<NodeId>, ) -> NodeId

Dense linear layer: matmul(input, weight) with optional rank-1 bias.

pub fn linear_bias_act( &mut self, input: NodeId, weight: NodeId, bias: Option<NodeId>, activation: Option<Activation>, ) -> NodeId

Dense linear with optional bias and epilogue activation.

pub fn linear_fused( &mut self, input: NodeId, weight: NodeId, bias: NodeId, activation: Option<Activation>, out_shape: Shape, ) -> NodeId

Emit Op::FusedMatMulBiasAct directly — deterministic fusion without relying on the FuseMatMulBiasAct pass.

pub fn shared_matmul_pair( &mut self, input: NodeId, w_first: NodeId, w_second: NodeId, ) -> (NodeId, NodeId)

Two matmuls sharing the same input — canonical gate+up / QKV pattern for FuseSharedInputMatMul.

Returns (first, second) in declaration order. For SwiGLU, pass up weight first and gate weight second so the post-concat narrow layout matches FuseSwiGLU (up @ 0, gate @ N).

pub fn swiglu_ffn( &mut self, input: NodeId, up_w: NodeId, gate_w: NodeId, down_w: NodeId, ) -> NodeId

SwiGLU FFN block: shared-input gate+up → silu(gate) * up → down proj.

Weight order matches FuseSwiGLU’s canonical narrow layout (up projection first, gate projection second).

pub fn fused_swiglu_ffn( &mut self, input: NodeId, up_w: NodeId, gate_w: NodeId, down_w: NodeId, out_shape: Shape, ) -> NodeId

Fully fused SwiGLU FFN: concat weights → single matmul → Op::FusedSwiGLU → down projection. Matches the rewrite performed by FuseSwiGLUDualMatmul without relying on the pass.

impl Graph

pub fn conv2d( &mut self, input: NodeId, weight: NodeId, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], groups: usize, ) -> NodeId

2D convolution on NCHW tensors (Op::Conv). Weight [C_out, C_in/g, kH, kW].

pub fn im2col( &mut self, input: NodeId, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ) -> NodeId

NCHW im2col (Op::Im2Col). Output [M, C·kH·kW].

pub fn conv_transpose2d( &mut self, input: NodeId, weight: NodeId, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], output_padding: [usize; 2], groups: usize, ) -> NodeId

2D transposed convolution on NCHW. Weight [C_in, C_out/g, kH, kW].

impl Graph

pub fn conv3d( &mut self, input: NodeId, weight: NodeId, stride: [usize; 3], padding: [usize; 3], dilation: [usize; 3], groups: usize, ) -> NodeId

3-D convolution on NCDHW tensors (Op::Conv3d).

• input — [N, C_in, D, H, W].
• weight — [C_out, C_in/groups, kD, kH, kW] (PyTorch Conv3d layout).

Kernel size is read from the weight. Returns [N, C_out, D_out, H_out, W_out] with X_out = floor((X + 2·p − dil·(K−1) − 1) / stride) + 1.

pub fn conv_transpose3d( &mut self, input: NodeId, weight: NodeId, stride: [usize; 3], padding: [usize; 3], dilation: [usize; 3], output_padding: [usize; 3], groups: usize, ) -> NodeId

3-D transposed convolution on NCDHW (Op::ConvTranspose3d), the learned upsampler in MONAI 3-D U-Net decoders.

• input — [N, C_in, D, H, W].
• weight — [C_in, C_out/groups, kD, kH, kW] (PyTorch ConvTranspose3d).

Returns [N, C_out, D_out, H_out, W_out] with X_out = (X−1)·stride − 2·p + dil·(K−1) + output_padding + 1.

pub fn interpolate3d( &mut self, x: NodeId, out_dhw: [usize; 3], mode: InterpMode, align_corners: bool, ) -> NodeId

Resample the spatial [D, H, W] axes of a [N, C, D, H, W] tensor to out_dhw, as a pure decomposition into transpose_/reshape_/mm (no new op). InterpMode::Linear gives separable trilinear; InterpMode::Nearest replicates the nearest source voxel.

align_corners selects the sampling grid (exposed so a decoder can match its reference):

• true — endpoint mapping pos = j·(L_in−1)/(L_out−1) (torch.nn.functional.interpolate(..., align_corners=True)).
• false — half-pixel mapping pos = (j+0.5)·L_in/L_out − 0.5, clamped to [0, L_in−1] — what MONAI UNETR’s input branch uses (trilinear, align_corners=False).

impl Graph

pub fn hilbert(&mut self, x: NodeId) -> (NodeId, NodeId)

Analytic signal of the last axis. Returns (real, imag), each the same shape as x; real ≈ x and imag = H(x) (the Hilbert transform).

pub fn envelope(&mut self, x: NodeId) -> NodeId

Amplitude envelope |x + i·H(x)| of the last axis (same shape as x).

pub fn instantaneous_phase(&mut self, x: NodeId) -> NodeId

Instantaneous phase atan2(H(x), x) of the last axis, in radians.

Uses the branch-cut-free identity atan2(y, x) = 2·atan(y / (√(x²+y²) + x)), exact everywhere except the negative real axis (a measure-zero set).

pub fn fir_filtfilt(&mut self, x: NodeId, taps: &[f32]) -> NodeId

Zero-phase FIR filtering: reverse(fir(reverse(fir(x)))).

taps are the FIR coefficients (host constant). Output has the same last-axis length as x; the two passes cancel the filter’s phase, so the effective magnitude response is |H(f)|². This is the zero-phase band-pass the exg preprocessor uses for MNE parity.

pub fn biquad(&mut self, x: NodeId, b: [f32; 3], a: [f32; 3]) -> NodeId

One biquad (2nd-order IIR) section applied along the last axis, Direct-Form II Transposed. b = [b0, b1, b2], a = [a0, a1, a2] (coefficients are normalized by a0). This is a true recurrence — y[n] = b0·x[n] + z1[n−1], etc. — evaluated with an Op::Scan over time, so it currently lowers on the CPU backend (where Op::Scan runs); GPU support tracks Op::Scan GPU lowering.

Any rank ≥ 1 is accepted; leading axes are independent channels.

pub fn sosfilt( &mut self, x: NodeId, sections: &[([f32; 3], [f32; 3])], ) -> NodeId

Cascade of second-order sections (scipy sosfilt): apply each (b, a) biquad in series along the last axis.

pub fn sosfiltfilt( &mut self, x: NodeId, sections: &[([f32; 3], [f32; 3])], ) -> NodeId

Zero-phase IIR filtering (scipy sosfiltfilt): forward pass, reverse, forward pass, reverse. Cancels the SOS cascade’s phase, leaving the squared magnitude response.

pub fn resample_poly( &mut self, x: NodeId, up: usize, down: usize, taps: &[f32], ) -> NodeId

Polyphase rational resampling of the last axis by up/down.

Zero-stuffs by up, applies the FIR anti-alias filter taps (scaled by up to preserve amplitude), then decimates by down. Output last-axis length is ceil(N·up / down). Composed from reshape/concat/FFT-conv/narrow, so it runs on every backend that supports those (all Apple-Silicon backends).

impl Graph

pub fn binary( &mut self, op: BinaryOp, lhs: NodeId, rhs: NodeId, out_shape: Shape, ) -> NodeId

Binary element-wise operation.

pub fn activation( &mut self, act: Activation, input: NodeId, shape: Shape, ) -> NodeId

Unary activation.

pub fn quantize(&mut self, x: NodeId, scale: f32, zero_point: i32) -> NodeId

Per-tensor INT8 quantization. Output dtype = I8, same shape otherwise. scale and zero_point apply uniformly to every element. Use quantize_per_channel when weights deserve per-channel scales (the standard PTQ improvement).

pub fn quantize_per_channel( &mut self, x: NodeId, axis: usize, scales: Vec<f32>, zero_points: Vec<i32>, ) -> NodeId

Per-channel INT8 quantization. scales and zero_points must each have length input.dim(axis); the kernel picks the i-th pair when quantizing the i-th slice along axis. The most common usage is axis = 0 for a [C_out, C_in, kH, kW] conv weight (one scale per output channel).

pub fn dequantize(&mut self, x: NodeId, scale: f32, zero_point: i32) -> NodeId

Per-tensor INT8 dequantization (inverse of quantize). Output dtype is f32.

pub fn dequantize_per_channel( &mut self, x: NodeId, axis: usize, scales: Vec<f32>, zero_points: Vec<i32>, ) -> NodeId

Per-channel INT8 dequantization (inverse of quantize_per_channel).

impl Graph

pub fn pad_last_axis_to_pow2(&mut self, x: NodeId) -> NodeId

Zero-pad the last axis to the next power of two (no-op when already pow2).

pub fn split_spectrum(&mut self, spectrum: NodeId) -> (NodeId, NodeId)

Split a 2N real-block spectrum into separate real / imag tensors.

pub fn fft_real(&mut self, x: NodeId, norm: FftNorm) -> (NodeId, NodeId)

Real-input FFT (gpu-fft fft): auto zero-pads to pow2, returns (re, im).

pub fn fft_batch_real(&mut self, x: NodeId, norm: FftNorm) -> (NodeId, NodeId)

Batched real-input FFT — same as fft_real when the last axis is signal length; leading axes are independent batch dimensions.

pub fn rfft(&mut self, x: NodeId, norm: FftNorm) -> (NodeId, NodeId)

Real-input FFT with half-spectrum output (n_pad/2 + 1 complex bins).

The input is zero-padded to the next power of two along the last axis before the transform, matching NumPy rfft padding semantics.

pub fn rfft_exact( &mut self, x: NodeId, n: usize, norm: FftNorm, ) -> (NodeId, NodeId)

Exact real-input FFT for an arbitrary length n (no power-of-two padding): returns the half-spectrum (re, im) with n/2 + 1 complex bins along the last axis.

Unlike Self::rfft — which zero-pads the last axis to next_pow2(n) and therefore samples the padded transform’s frequencies — this computes the genuine n-point real DFT via a constant DFT-matrix matmul, exact for every n (odd, prime, non-pow2). It is a pure decomposition over existing graph ops (Op::Constant, Op::MatMul, elementwise), so it lowers on every backend and stays differentiable — the same strategy as Self::interpolate1d. Two baked [n, n/2+1] matrices give

  re[k] =  Σ_t x[t]·cos(2π k t / n)          COS [n, n/2+1]
  im[k] = −Σ_t x[t]·sin(2π k t / n)          NSIN[n, n/2+1]  (= −sin)

matching the sign convention of torch.fft.rfft and this crate’s radix-2 Self::rfft (forward transform X[k] = Σ x[t] e^{−2πi k t / n}). The last axis of x must be the real signal of length n; leading axes are independent batch dimensions. norm scales the forward transform exactly as Self::rfft (FftNorm::output_scale(n, false): Backward/Forward → 1, Ortho → 1/√n). PyTorch norm='forward'’s 1/n forward scale is not an FftNorm variant — apply it explicitly on the result if needed (see rlx-cbramod’s spectral front-end).

Cost is O(n²) per row; for large n a Bluestein / mixed-radix kernel (the rlx-fft butterfly/Stockham path could host one) would be asymptotically cheaper. For the small EEG-tokenizer windows this targets (n = 200, 400) the exact matmul is both simpler and cheap enough.

pub fn rfft_exact_mag(&mut self, x: NodeId, n: usize, norm: FftNorm) -> NodeId

Magnitude of the exact arbitrary-n half-spectrum: sqrt(re² + im²), shape [.., n/2 + 1]. The imaginary sign is irrelevant to the magnitude, so this is the spectral front-end the EEG tokenizers (CBraMod, BrainBERT) want. norm scales as in Self::rfft_exact.

pub fn irfft( &mut self, re_half: NodeId, im_half: NodeId, n: usize, norm: FftNorm, ) -> NodeId

Inverse real FFT from half-spectrum (re, im) with Hermitian symmetry.

Mirrors the conjugate half of the spectrum (excluding DC and Nyquist) before calling Self::ifft_spectrum, then truncates to length n.

pub fn stft( &mut self, x: NodeId, frame_len: usize, hop: usize, norm: FftNorm, ) -> NodeId

Short-time Fourier transform: [..., T] → [frames, ..., 2·half] (re/im block per frame).

Each frame is rfft’d with length frame_len and hop hop along the last axis.

pub fn fft_conv1d( &mut self, a: NodeId, b: NodeId, n_fft: usize, norm: FftNorm, ) -> NodeId

1D convolution via the convolution theorem (rfft → complex multiply → irfft).

Both inputs are zero-padded to at least n_fft (or the next power of two covering len(a) + len(b) - 1 when n_fft is small).

pub fn fftfreq_tensor(&mut self, n: usize) -> NodeId

Constant tensor of FFT sample frequencies (length n, f64).

pub fn rfftfreq_tensor(&mut self, n: usize) -> NodeId

Constant tensor of rFFT sample frequencies (length n/2 + 1, f64).

pub fn psd_real(&mut self, x: NodeId, norm: FftNorm) -> NodeId

Power spectral density from real input: rfft → psd.

pub fn ifft_spectrum(&mut self, re: NodeId, im: NodeId, norm: FftNorm) -> NodeId

Inverse FFT from separate real / imag spectra (gpu-fft ifft real part).

pub fn psd(&mut self, re: NodeId, im: NodeId) -> NodeId

Power spectral density: (re² + im²) / N (gpu-fft psd::psd).

impl Graph

pub fn input(&mut self, name: impl Into<String>, shape: Shape) -> NodeId

Graph input (runtime-provided tensor).

pub fn param(&mut self, name: impl Into<String>, shape: Shape) -> NodeId

Model parameter (weight loaded at init).

pub fn add_node(&mut self, op: Op, inputs: Vec<NodeId>, shape: Shape) -> NodeId

Generic node constructor for custom ops.

pub fn custom_op( &mut self, name: impl Into<String>, attrs: Vec<u8>, inputs: Vec<NodeId>, ) -> NodeId

Build an Op::Custom node, dispatching shape inference through the global op registry. The named op must already be registered via crate::register_op; attrs is forwarded verbatim to the impl’s infer_shape (and later, at execution time, to its per-backend kernel).

Panics if name is not registered or if inputs.len() does not match the registered num_inputs() — both are programmer errors that should fail loudly at graph-build time, not silently at execution.

pub fn custom_op_packed( &mut self, name: impl Into<String>, attrs: Vec<u8>, inputs: Vec<NodeId>, out_shape: Shape, ) -> NodeId

Build an Op::Custom node with a caller-supplied output shape, bypassing the registry’s infer_shape. Use this for ops whose output shape can’t be determined by static input shapes alone — most importantly, ops with multiple logical outputs packed into one buffer.

The canonical multi-output pattern:

// Sparse-LU returns L_values + U_values packed end-to-end.
// Caller knows nnz_L and nnz_U from the symbolic factor.
let lu = g.custom_op_packed(
    "sparse_lu",
    attrs,
    vec![A, b],
    Shape::new(&[nnz_L + nnz_U], DType::F64),
);
let l_vals = g.narrow_(lu, 0, 0, nnz_L);
let u_vals = g.narrow_(lu, 0, nnz_L, nnz_U);

The op must still be registered (so num_inputs validation and autodiff routing still work); only the shape is overridden.

pub fn fft(&mut self, x: NodeId, inverse: bool) -> NodeId

1D FFT along the last axis.

• F32 / F64 — 2N real-block layout: last axis is [re…, im…].
• C64 — interleaved [re, im] pairs per complex element.

Output shape matches input. Radix-2 when N is a power of two, Bluestein otherwise. Default normalization is unnormalized (FftNorm::Backward; ifft(fft(x)) = N·x).

pub fn fft_norm(&mut self, x: NodeId, inverse: bool, norm: FftNorm) -> NodeId

1D FFT with explicit normalization mode.

pub fn fft_butterfly_stage( &mut self, state: NodeId, gate: NodeId, rev: NodeId, tw_re: NodeId, tw_im: NodeId, stage: u32, n_fft: u32, ) -> NodeId

Ternary pruned radix-2 butterfly stage — see Op::FftButterflyStage.

pub fn fft_axis(&mut self, x: NodeId, axis: usize, inverse: bool) -> NodeId

1D FFT along an arbitrary axis. Lowers to Transpose(axis ↔ last) → Fft(last) → Transpose(last ↔ axis).

AD is free: both Op::Transpose and Op::Fft have VJP/JVP rules.

pub fn fftn(&mut self, x: NodeId, axes: &[usize], inverse: bool) -> NodeId

N-dimensional FFT along axes (NumPy fftn semantics).

Applies a 1D FFT along each listed axis in ascending order. Empty axes is a no-op. For multi-axis transforms on tensors with more than one spatial dimension, use DType::C64; the F32/F64 2N-block layout only describes a single complex axis.

pub fn ifftn(&mut self, x: NodeId, axes: &[usize]) -> NodeId

Inverse N-dimensional FFT — alias for fftn(..., inverse: true).

impl Graph

pub fn matmul(&mut self, lhs: NodeId, rhs: NodeId, out_shape: Shape) -> NodeId

Matrix multiply.

pub fn dense_solve(&mut self, a: NodeId, b: NodeId, out_shape: Shape) -> NodeId

Dense linear solve x = A⁻¹·b. A must be [N, N]; b is [N] for a single right-hand side or [N, K] for multiple. out_shape matches b’s shape.

pub fn batched_dense_solve( &mut self, a: NodeId, b: NodeId, out_shape: Shape, ) -> NodeId

Batched dense linear solve. A is [B, N, N]; b is [B, N] (single-RHS) or [B, N, K] (multi-RHS). Per-batch independent — each slice solved as a separate dense_solve. Typically constructed by vmap of dense_solve.

pub fn lora_matmul( &mut self, x: NodeId, w: NodeId, a: NodeId, b: NodeId, scale: f32, shape: Shape, ) -> NodeId

Fused LoRA matmul: out = x·W + scale * (x·A)·B. Inputs: x [m, k], w [k, n], a [k, r], b [r, n]. r is the LoRA rank; scale is the alpha/rank coefficient.

pub fn dequant_matmul( &mut self, x: NodeId, w_q: NodeId, scale: NodeId, zp: NodeId, scheme: QuantScheme, shape: Shape, ) -> NodeId

Fused dequant + matmul. See Op::DequantMatMul for per-scheme input layout (4 inputs for legacy/NVFP4, 2 for GGUF).

pub fn dequant_matmul_packed( &mut self, x: NodeId, packed_w: NodeId, scheme: QuantScheme, shape: Shape, ) -> NodeId

GGUF / K-quant packed weights — [x, packed_w_bytes] only.

pub fn dequant_matmul_nvfp4( &mut self, x: NodeId, w_q: NodeId, block_scales: NodeId, global_scale: NodeId, shape: Shape, ) -> NodeId

NVFP4 (E2M1) block matmul — group size 16, FP8 block scales, optional f32 global scale (defaults to 1.0 when unset at runtime).

pub fn fused_matmul_bias_act( &mut self, input: NodeId, weight: NodeId, bias: NodeId, activation: Option<Activation>, shape: Shape, ) -> NodeId

Fused matmul + bias + activation (created by optimization passes).

pub fn q_matmul( &mut self, x: NodeId, w: NodeId, bias: NodeId, x_zp: i32, w_zp: i32, out_zp: i32, mult: f32, out_shape: Shape, ) -> NodeId

Real INT8-arithmetic matmul: i8 inputs, i32 bias, i8 output. mult = x_scale · w_scale / out_scale. Caller’s responsible for asserting the input dtypes — the builder just plumbs the shape with dtype = I8 since that’s what the kernel writes.

pub fn q_conv2d( &mut self, x: NodeId, w: NodeId, bias: NodeId, kernel_size: Vec<usize>, stride: Vec<usize>, padding: Vec<usize>, dilation: Vec<usize>, groups: usize, x_zp: i32, w_zp: i32, out_zp: i32, mult: f32, out_shape: Shape, ) -> NodeId

Real INT8-arithmetic 2-D convolution. NCHW layout matching Op::Conv. mult = x_scale · w_scale / out_scale.

impl Graph

pub fn layer_norm2d( &mut self, input: NodeId, gamma: NodeId, beta: NodeId, eps: f32, ) -> NodeId

LayerNorm2d on NCHW (normalize across channels at each spatial position).

pub fn group_norm( &mut self, input: NodeId, gamma: NodeId, beta: NodeId, num_groups: usize, eps: f32, ) -> NodeId

Group normalization on NCHW.

pub fn batch_norm_inference( &mut self, input: NodeId, gamma: NodeId, beta: NodeId, running_mean: NodeId, running_var: NodeId, eps: f32, ) -> NodeId

BatchNorm inference (frozen running mean/variance).

pub fn layer_norm( &mut self, input: NodeId, gamma: NodeId, beta: NodeId, axis: i32, eps: f32, shape: Shape, ) -> NodeId

Layer normalization.

pub fn fused_residual_ln( &mut self, x: NodeId, residual: NodeId, bias: Option<NodeId>, gamma: NodeId, beta: NodeId, eps: f32, shape: Shape, ) -> NodeId

Fused residual + bias + layer norm (created by optimization passes).

pub fn fused_residual_rms_norm( &mut self, x: NodeId, residual: NodeId, bias: Option<NodeId>, gamma: NodeId, beta: NodeId, eps: f32, shape: Shape, ) -> NodeId

Fused residual + bias + RMS norm (created by optimization passes).

impl Graph

pub fn reduce( &mut self, input: NodeId, op: ReduceOp, axes: Vec<usize>, keep_dim: bool, shape: Shape, ) -> NodeId

Reduce.

pub fn softmax(&mut self, input: NodeId, axis: i32, shape: Shape) -> NodeId

Softmax.

pub fn cumsum( &mut self, input: NodeId, axis: i32, exclusive: bool, shape: Shape, ) -> NodeId

Cumulative sum along an axis (output shape == input shape).

pub fn argmax( &mut self, input: NodeId, axis: usize, keep_dim: bool, shape: Shape, ) -> NodeId

Index of the max along axis (f32-encoded indices).

pub fn argmin( &mut self, input: NodeId, axis: usize, keep_dim: bool, shape: Shape, ) -> NodeId

Index of the min along axis (f32-encoded indices).

pub fn sample( &mut self, logits: NodeId, top_k: usize, top_p: f32, temperature: f32, seed: u64, output_shape: Shape, ) -> NodeId

Fused sample: logits → token id (one f32-encoded id per row). output_shape should be [batch] (one id per logit row).

impl Graph

pub fn reshape( &mut self, input: NodeId, new_shape: Vec<i64>, out_shape: Shape, ) -> NodeId

Reshape.

pub fn gather( &mut self, table: NodeId, indices: NodeId, axis: usize, shape: Shape, ) -> NodeId

Gather (embedding lookup).

pub fn concat( &mut self, inputs: Vec<NodeId>, axis: usize, shape: Shape, ) -> NodeId

Concatenate tensors along an axis.

pub fn reverse(&mut self, input: NodeId, axes: Vec<usize>) -> NodeId

Reverse (flip) element order along each axis in axes. Output shape is the same as the input; only the listed axes flip (batch-general).

impl Graph

pub fn selective_scan( &mut self, x: NodeId, delta: NodeId, a: NodeId, b: NodeId, c: NodeId, state_size: usize, shape: Shape, ) -> NodeId

Mamba-style selective scan: y = SSM(x, Δ, A, B, C). Inputs: x [b,s,h], delta [b,s,h], a [h,n], b [b,s,n], c [b,s,n]. Output [b,s,h]. n is the state size.

pub fn gated_delta_net( &mut self, q: NodeId, k: NodeId, v: NodeId, g: NodeId, beta: NodeId, state_size: usize, shape: Shape, ) -> NodeId

Gated DeltaNet linear-attention scan (Qwen3.5/3.6 trunk, Qwen3-Next, Kimi-Linear). See Op::GatedDeltaNet for the recurrence math. All five inputs are f32. Shapes: q,k,v: [b, s, h_v, n]; g,beta: [b, s, h_v]. Output: [b, s, h_v, n]. State is implicit (reset per batch) unless carry_state is set — then pass state as a sixth input.

pub fn gated_delta_net_carry( &mut self, q: NodeId, k: NodeId, v: NodeId, g: NodeId, beta: NodeId, state: NodeId, state_size: usize, shape: Shape, ) -> NodeId

Same as Self::gated_delta_net but threads state [b, h_v, n, n] in/out for decode-mode recurrence.

pub fn lstm( &mut self, x: NodeId, w_ih: NodeId, w_hh: NodeId, bias: NodeId, hidden_size: usize, num_layers: usize, bidirectional: bool, shape: Shape, ) -> NodeId

Multi-layer (optionally bidirectional) LSTM with packed weights. See Op::Lstm for the gate math, weight packing, and shapes. Output shape = [batch, seq, D*hidden] (D = 2 iff bidirectional). Initial h0/c0 are zero.

pub fn lstm_carry( &mut self, x: NodeId, w_ih: NodeId, w_hh: NodeId, bias: NodeId, h0: NodeId, c0: NodeId, hidden_size: usize, num_layers: usize, bidirectional: bool, shape: Shape, ) -> NodeId

Same as Self::lstm but threads decode state: h0/c0 [L*D, batch, hidden] in, hn/cn written back in place.

pub fn gru( &mut self, x: NodeId, w_ih: NodeId, w_hh: NodeId, b_ih: NodeId, b_hh: NodeId, h0: Option<NodeId>, hidden_size: usize, num_layers: usize, bidirectional: bool, shape: Shape, ) -> NodeId

Multi-layer (optionally bidirectional) GRU. Inputs [x, w_ih, w_hh, b_ih, b_hh] (+ [h0] when carry). See Op::Gru. Output shape = [batch, seq, D*hidden].

pub fn rnn( &mut self, x: NodeId, w_ih: NodeId, w_hh: NodeId, bias: NodeId, h0: Option<NodeId>, hidden_size: usize, num_layers: usize, bidirectional: bool, relu: bool, shape: Shape, ) -> NodeId

Multi-layer (optionally bidirectional) Elman RNN. Inputs [x, w_ih, w_hh, bias] (+ [h0] when carry). relu selects the activation (else tanh). See Op::Rnn. Output [batch, seq, D*hidden].

pub fn mamba2( &mut self, x: NodeId, dt: NodeId, a: NodeId, b: NodeId, c: NodeId, head_dim: usize, state_size: usize, shape: Shape, ) -> NodeId

Mamba-2 / SSD scalar-decay SSM scan. See Op::Mamba2. Inputs [x, dt, a, b, c]; output shape = x shape [B,S,H,P].

pub fn scan(&mut self, init: NodeId, body: Graph, length: u32) -> NodeId

Bounded scan returning the final carry. Body must have exactly one Op::Input (the carry) and one output, both same shape as init. Output shape matches init.

pub fn scan_checkpointed( &mut self, init: NodeId, body: Graph, length: u32, num_checkpoints: u32, ) -> NodeId

Bounded scan with recursive checkpointing for memory-bounded backward AD. Equivalent to Self::scan for the forward computation, but during backward only num_checkpoints carry values are cached; intermediate carries are recomputed via the body. Memory: O(num_checkpoints · carry_size). Time: forward unchanged; backward O(length) (segment-cached).

The AD pre-pass propagates num_checkpoints into the rewritten trajectory-saving Scan and into the emitted ScanBackward, so a single call to crate::Graph::scan_checkpointed is enough to enable the memory bound across the whole forward+backward pipeline.

pub fn scan_with_bcasts_and_xs( &mut self, init: NodeId, bcasts: &[NodeId], xs: &[NodeId], body: Graph, length: u32, ) -> NodeId

Bounded scan with broadcast and per-step inputs.

Body Op::Inputs in NodeId order: [carry, bcast_0..bcast_{B-1}, x_t_0..x_t_{X-1}]. Bcast inputs keep their natural shape (the CPU executor fills them once before the scan loop). xs[i] has shape [length, *per_step] and the body sees xs[i][t] per iteration. Output shape matches init.

pub fn scan_with_xs( &mut self, init: NodeId, xs: &[NodeId], body: Graph, length: u32, ) -> NodeId

Bounded scan with per-step xs inputs returning the final carry. Body has 1 + xs.len() Op::Inputs in NodeId construction order (first declared is the carry; the remaining match xs in order). Each xs[i] has shape [length, *per_step_shape_i]; the body sees a per_step_shape_i slice on iteration t.

pub fn scan_backward( &mut self, init: NodeId, trajectory: NodeId, upstream: NodeId, xs: &[NodeId], body_vjp: Graph, length: u32, save_trajectory: bool, out_shape: Shape, ) -> NodeId

Reverse-mode AD companion to Self::scan / Self::scan_trajectory. Typically constructed by the autodiff pass, not by hand.

xs is the list of per-step input tensors (must match the forward Op::Scan’s xs in count, order, and per-step shape). Body_vjp’s 1 + xs.len() + 1 Op::Inputs match the forward body’s inputs plus a fresh "d_output" Input.

pub fn scan_backward_with_checkpoints( &mut self, init: NodeId, trajectory: NodeId, upstream: NodeId, xs: &[NodeId], body_vjp: Graph, length: u32, save_trajectory: bool, num_checkpoints: u32, forward_body: Option<Graph>, out_shape: Shape, ) -> NodeId

Lower-level scan_backward with explicit checkpointing config. num_checkpoints == 0 (default) means no checkpointing — the trajectory cache holds every step’s carry. 0 < K < length enables segment-cached recompute via forward_body (must be Some).

pub fn scan_backward_xs( &mut self, init: NodeId, trajectory: NodeId, upstream: NodeId, xs: &[NodeId], body_vjp: Graph, length: u32, save_trajectory: bool, xs_idx: u32, out_shape: Shape, ) -> NodeId

Per-step xs gradient companion to Self::scan_backward. Same inputs and same body_vjp graph, plus an xs_idx selecting which body_vjp output to stack into the result. Output shape is [length, *per_step_xs_shape].

pub fn scan_backward_xs_with_checkpoints( &mut self, init: NodeId, trajectory: NodeId, upstream: NodeId, xs: &[NodeId], body_vjp: Graph, length: u32, save_trajectory: bool, xs_idx: u32, num_checkpoints: u32, forward_body: Option<Graph>, out_shape: Shape, ) -> NodeId

pub fn custom_fn( &mut self, inputs: Vec<NodeId>, fwd_body: Graph, vjp_body: Option<Graph>, jvp_body: Option<Graph>, ) -> NodeId

User-defined sub-graph with optional override AD rules. JAX-shaped custom_vjp / custom_jvp — see Op::CustomFn.

inputs.len() must equal the number of Op::Input nodes in fwd_body. Output shape is inferred from fwd_body’s declared output. When supplied, vjp_body and jvp_body must follow the conventions documented on Op::CustomFn (special-named "primal_output" / "d_output" / "tangent_*" Inputs).

pub fn custom_fn_multi( &mut self, inputs: Vec<NodeId>, fwd_body: Graph, ) -> MultiOutputHandle

Multi-output custom_fn via the concat-with-Narrow design: rewrites fwd_body to flatten + concat its K declared outputs into a single 1-D F32 output, wraps that as Op::CustomFn, and returns a MultiOutputHandle the caller uses to extract each sub-output via Op::Narrow + Op::Reshape.

Per PLAN line 484, this avoids rewriting rlx’s “1 Op = 1 output” IR contract: the wrapped Op::CustomFn still has one output (the flat concat), and MultiOutputHandle::output(g, i) materializes component i lazily on the outer graph.

Constraints (MVP):

• All sub-outputs must be DType::F32. Tuples-of-mixed-dtype need either a per-dtype split or a future tuple-type extension.
• All sub-output shapes must be statically known (no Dim::Dynamic).
• vjp_body / jvp_body aren’t yet rewritten through the concat — caller must provide bodies that already expect the flat-concat output convention if they need custom AD.

pub fn scan_trajectory( &mut self, init: NodeId, body: Graph, length: u32, ) -> NodeId

Bounded scan returning the stacked trajectory. Output shape is [length, *init.shape] — row t is the carry after step t+1, so row length-1 equals the result of plain Self::scan.

impl Graph

pub fn spectrogram( &mut self, x: NodeId, frame_len: usize, hop: usize, window: WindowKind, power: bool, log: bool, ) -> NodeId

Windowed short-time spectrogram of the last axis.

[..., T] → [n_frames, ..., n_bins] where n_frames = 1 + (T − frame_len)/hop and n_bins = next_pow2(frame_len)/2 + 1.

• power — return |X|² (true) or magnitude |X| (false).
• log — apply log(· + 1e-8) to the result (log-power / log-magnitude).

pub fn band_power( &mut self, x: NodeId, sample_rate: f32, bands: &[(f32, f32)], ) -> NodeId

Per-band power of the last axis via a single rfft.

[..., T] → [..., n_bands]. Each band (lo, hi) in Hz sums the (unnormalized) power |X|² over the rFFT bins whose center frequency falls in [lo, hi], using sample_rate to map Hz → bin index.

pub fn differential_entropy( &mut self, x: NodeId, sample_rate: f32, bands: &[(f32, f32)], ) -> NodeId

Per-band differential entropy: 0.5·log(2πe · band_power + 1e-8).

For a band-limited Gaussian signal this equals its differential entropy up to the constant — the classic EEG “DE” feature (SEED, MAET, FoME). Returns [..., n_bands].

impl Graph

pub fn gaussian_splat_render( &mut self, inputs: GaussianSplatInputs, params: GaussianSplatRenderParams, ) -> NodeId

First-class CPU reference Gaussian splat forward render.

See Op::GaussianSplatRender for the seven-input contract and GaussianSplatRenderParams for framebuffer settings.

pub fn gaussian_splat_render_meta( &mut self, camera_position: [f32; 3], camera_target: [f32; 3], camera_up: [f32; 3], fov_y_degrees: f32, near: f32, far: f32, background: [f32; 3], params: GaussianSplatRenderParams, ) -> NodeId

Build the 23-float meta vector expected by Op::GaussianSplatRender.

pub fn gaussian_splat_prepare( &mut self, inputs: GaussianSplatInputs, params: GaussianSplatRenderParams, ) -> NodeId

Strict IR stage 1: project + bin + sort + rays → packed prepare buffer.

pub fn gaussian_splat_rasterize( &mut self, prep: NodeId, meta: NodeId, params: GaussianSplatRenderParams, ) -> NodeId

Strict IR stage 2: rasterize from prepare buffer + meta.

pub fn gaussian_splat_render_decomposed( &mut self, inputs: GaussianSplatInputs, params: GaussianSplatRenderParams, ) -> NodeId

Decomposed strict-IR forward: prepare → rasterize.

pub fn gaussian_splat_render_backward( &mut self, inputs: GaussianSplatInputs, d_loss_rgba: NodeId, params: GaussianSplatBackwardParams, ) -> NodeId

Backward pass for Op::GaussianSplatRender (packed scene gradients).

impl Graph

pub fn conv_transpose1d( &mut self, input: NodeId, weight: NodeId, kernel: usize, stride: usize, padding: usize, dilation: usize, output_padding: usize, groups: usize, ) -> NodeId

Transposed 1-D convolution (a.k.a. deconvolution / fractionally-strided conv), the learned upsampler in Conv1d U-Net decoders.

• input — [N, C_in, L].
• weight — [C_in, C_out/groups, K] (PyTorch ConvTranspose1d layout).

Returns [N, C_out, L_out] with L_out = (L−1)·stride − 2·padding + dilation·(K−1) + output_padding + 1.

pub fn interpolate1d( &mut self, x: NodeId, l_out: usize, mode: InterpMode, ) -> NodeId

Resample the last axis of x from L_in to l_out samples.

Works for any rank ≥ 1; leading axes are treated as batch. Uses align-corners endpoint mapping (pos = j·(L_in−1)/(l_out−1)), matching torch.nn.functional.interpolate(..., align_corners=True).

impl Graph

pub fn vector_quantize( &mut self, x: NodeId, codebook: NodeId, metric: VqMetric, ) -> (NodeId, NodeId)

Nearest-codebook assignment.

• x — inputs [N, D].
• codebook — codes [K, D].

Returns (indices, quantized) where indices is [N] (f32-encoded code ids, ready to feed gather) and quantized is [N, D], the selected code vectors. The straight-through estimator for training is left to the caller (x + stop_gradient(quantized − x)), matching the usual VQ-VAE recipe.

pub fn residual_vq( &mut self, x: NodeId, codebooks: &[NodeId], metric: VqMetric, ) -> (Vec<NodeId>, NodeId)

Residual (multi-stage) vector quantization.

Quantizes x against codebooks[0], subtracts the chosen code, then quantizes the residual against codebooks[1], and so on. Returns the per-level indices (one [N] tensor per stage) and the summed reconstruction quantized[N, D] (Σ_level code_level). This is the RVQ tokenizer in NeuroRVQ / BrainRVQ.

Trait Implementations

impl Clone for Graph

fn clone(&self) -> Graph

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for Graph

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

Formats the value using the given formatter.

impl<'de> Deserialize<'de> for Graph

fn deserialize<__D>( __deserializer: __D, ) -> Result<Graph, <__D as Deserializer<'de>>::Error>

where __D: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl Display for Graph

Pretty-print the graph in a readable IR format.

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

Formats the value using the given formatter.

impl From<MirModule> for Graph

fn from(mir: MirModule) -> Graph

Converts to this type from the input type.

impl GraphExt for Graph

fn mm(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId

fn add(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId

fn sub(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId

fn mul(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId

fn div(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId

fn gelu(&mut self, x: NodeId) -> NodeId

fn gelu_approx(&mut self, x: NodeId) -> NodeId

Tanh-approximation GELU (PyTorch’s default gelu formula, also candle’s Tensor::gelu). Use this when porting models whose reference implementations use the tanh form for numerical parity (e.g. DINOv2, many ViTs).

fn silu(&mut self, x: NodeId) -> NodeId

fn relu(&mut self, x: NodeId) -> NodeId

fn exp(&mut self, x: NodeId) -> NodeId

fn sqrt(&mut self, x: NodeId) -> NodeId

fn neg(&mut self, x: NodeId) -> NodeId

fn tanh(&mut self, x: NodeId) -> NodeId

fn ln(&mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId

fn layer_norm2d( &mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32, ) -> NodeId

fn group_norm( &mut self, x: NodeId, gamma: NodeId, beta: NodeId, num_groups: usize, eps: f32, ) -> NodeId

fn conv2d( &mut self, input: NodeId, weight: NodeId, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], groups: usize, ) -> NodeId

fn conv_transpose2d( &mut self, input: NodeId, weight: NodeId, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], output_padding: [usize; 2], groups: usize, ) -> NodeId

fn rms_norm( &mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32, ) -> NodeId

fn sum(&mut self, x: NodeId, axes: Vec<usize>, keep_dim: bool) -> NodeId

fn mean(&mut self, x: NodeId, axes: Vec<usize>, keep_dim: bool) -> NodeId

fn sm(&mut self, x: NodeId, axis: i32) -> NodeId

fn reshape_(&mut self, x: NodeId, new_shape: Vec<i64>) -> NodeId

fn transpose_(&mut self, x: NodeId, perm: Vec<usize>) -> NodeId

fn narrow_( &mut self, x: NodeId, axis: usize, start: usize, len: usize, ) -> NodeId

fn concat_(&mut self, inputs: Vec<NodeId>, axis: usize) -> NodeId

fn gather_(&mut self, table: NodeId, indices: NodeId, axis: usize) -> NodeId

fn eq(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId

fn lt(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId

fn attention_( &mut self, q: NodeId, k: NodeId, v: NodeId, mask: NodeId, num_heads: usize, head_dim: usize, ) -> NodeId

fn rope( &mut self, x: NodeId, cos: NodeId, sin: NodeId, head_dim: usize, ) -> NodeId

fn rope_n( &mut self, x: NodeId, cos: NodeId, sin: NodeId, head_dim: usize, n_rot: usize, ) -> NodeId

Partial RoPE: rotate the first n_rot dims (NeoX offset n_rot/2).

fn rope_styled( &mut self, x: NodeId, cos: NodeId, sin: NodeId, head_dim: usize, style: RopeStyle, ) -> NodeId

RoPE with an explicit pairing flavor (crate::op::RopeStyle).

fn rope_n_styled( &mut self, x: NodeId, cos: NodeId, sin: NodeId, head_dim: usize, n_rot: usize, style: RopeStyle, ) -> NodeId

Partial RoPE with an explicit pairing flavor.

fn cast(&mut self, x: NodeId, to: DType) -> NodeId

fn try_constant(&mut self, value: f64, dtype: DType) -> Result<NodeId, String>

Fallible variant of GraphExt::constant. Returns an error when value is out of range for dtype or when dtype cannot be encoded directly (callers may lower f16 / bf16 via try_constant on F32 plus cast).

fn constant(&mut self, value: f64, dtype: DType) -> NodeId

Rank-0 broadcastable scalar (Op::Constant). f16 / bf16 are lowered as f32 constant + cast.

fn stop_gradient(&mut self, x: NodeId) -> NodeId

Identity forward, zero backward. The reverse-mode autodiff rule for Op::StopGradient returns no gradient contribution to the input. Equivalent to PyTorch’s tensor.detach() / jax.lax.stop_gradient / TF’s tf.stop_gradient.

impl PartialEq for Graph

fn eq(&self, other: &Graph) -> bool

Tests for self and other values to be equal, and is used by ==.

fn ne(&self, other: &Rhs) -> bool

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.

impl Serialize for Graph

fn serialize<__S>( &self, __serializer: __S, ) -> Result<<__S as Serializer>::Ok, <__S as Serializer>::Error>

where __S: Serializer,

Serialize this value into the given Serde serializer.

impl TryFrom<GraphModule> for Graph

type Error = LowerError

The type returned in the event of a conversion error.

fn try_from(module: GraphModule) -> Result<Graph, LowerError>

Performs the conversion.