rlx_tensor/

tensor.rs

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Symbolic tensor handle — zero payload bytes until compile + run.

// The binary/compare builder methods take `impl Into<BinaryRhs>` so callers can
// pass `&Tensor` or a scalar literal. `BinaryRhs` / `Scalar` are sealed
// conversion targets — deliberately reachable only through `Into`, never named
// by callers — so we accept them appearing in these public bounds/signatures.
// (`impl Into<BinaryRhs>` is a *bound*, hence `private_bounds`.)
#![allow(private_bounds, private_interfaces)]

use std::ops::{Add, Div, Mul, Neg, Sub};

use rlx_ir::op::{Activation, BinaryOp, CmpOp, MaskKind, ReduceOp};
use rlx_ir::{DType, Dim, GraphExt, NodeId, Op, ScaleLayout, ScaledFormat, Shape};

use crate::handle::GraphHandle;
use crate::scalar::{Scalar, promote_scalar};
use crate::slice::SliceSpec;

/// Generate same-shape unary [`Activation`] methods (`name => Variant`).
/// Each delegates to [`Tensor::map_act`], so adding a JAX-style elementwise
/// op is a one-line table entry, not a copied method body.
macro_rules! unary_ops {
    ($( $(#[$doc:meta])* $name:ident => $variant:ident ),+ $(,)?) => {
        impl Tensor {
            $(
                $(#[$doc])*
                pub fn $name(&self) -> Self {
                    self.map_act(Activation::$variant)
                }
            )+
        }
    };
}

/// Generate elementwise comparison methods (`name => Variant`) → `Bool` tensor.
/// Each accepts a `&Tensor` or a scalar (`x.gt(0.0)`), so masks read naturally
/// (both directly and from the `rlx!` DSL's `x > 0.0`).
macro_rules! cmp_ops {
    ($( $(#[$doc:meta])* $name:ident => $variant:ident ),+ $(,)?) => {
        impl Tensor {
            $(
                $(#[$doc])*
                pub fn $name<'r>(&self, rhs: impl Into<BinaryRhs<'r>>) -> Self {
                    self.map_cmp(rhs.into(), CmpOp::$variant)
                }
            )+
        }
    };
}

/// Generate axis-reduction methods (`name => Variant`).
macro_rules! reduce_ops {
    ($( $(#[$doc:meta])* $name:ident => $variant:ident ),+ $(,)?) => {
        impl Tensor {
            $(
                $(#[$doc])*
                pub fn $name(&self, axes: impl Into<Vec<usize>>, keep_dim: bool) -> Self {
                    self.map_reduce(ReduceOp::$variant, axes.into(), keep_dim)
                }
            )+
        }
    };
}

/// Convert a [`Dim`] to a reshape-style `i64` (dynamic → `-1`).
fn dim_to_i64(d: &Dim) -> i64 {
    match d {
        Dim::Static(n) => *n as i64,
        Dim::Dynamic(_) => -1,
    }
}

/// Render row-major `data` as nested brackets matching `dims`.
#[cfg(feature = "eval")]
fn fmt_nested(data: &[f32], dims: &[usize]) -> String {
    match dims {
        [] => format!("{}", data.first().copied().unwrap_or(0.0)),
        [_] => {
            let body = data
                .iter()
                .map(|x| format!("{x}"))
                .collect::<Vec<_>>()
                .join(", ");
            format!("[{body}]")
        }
        [outer, rest @ ..] => {
            let inner: usize = rest.iter().product();
            let parts = (0..*outer)
                .map(|i| fmt_nested(&data[i * inner..(i + 1) * inner], rest))
                .collect::<Vec<_>>()
                .join(", ");
            format!("[{parts}]")
        }
    }
}

#[cfg(feature = "eval")]
impl std::fmt::Display for Tensor {
    /// Realizes the tensor (compile + run) and prints values — the lazy graph
    /// is forced here, like ndarray's eager `Display`.
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.show())
    }
}

/// A symbolic tensor in a traced graph.
///
/// **Lazy by construction.** Despite the eager-looking operator syntax, a
/// `Tensor` is a deferred handle: `&a + &b`, `.relu()`, `.matmul(..)` etc.
/// only append IR nodes — no host or device buffers are allocated and nothing
/// executes. The graph is compiled (fused + memory-planned) and run **only**
/// when you force it with [`to_vec`](Self::to_vec) / [`realize`](Self::realize)
/// (both require the `eval` feature). This is what lets the whole expression
/// fuse instead of materializing every intermediate.
///
/// **Zero-copy clone & slice.** `clone()` just bumps a refcount on the shared
/// graph handle — no data is duplicated. `slice`/`narrow`/`reshape`/`select`
/// add a *view* node into the same graph (referencing the parent by id), so
/// they move zero payload bytes either. Use [`shares_graph`](Self::shares_graph)
/// to observe storage sharing and [`storage_bytes`](Self::storage_bytes) to
/// confirm no data was copied. (Combining two *independently-constructed*
/// tensors is the one case that copies — the right operand's nodes are merged
/// in once; see the `handle` module.)
#[derive(Clone, Debug)]
pub struct Tensor {
    pub(crate) handle: GraphHandle,
    pub(crate) id: NodeId,
}

impl Tensor {
    pub(crate) fn new(handle: GraphHandle, id: NodeId) -> Self {
        Self { handle, id }
    }

    /// IR node id (mix with raw [`rlx_ir::Graph`] builders via [`NodeId`]).
    pub fn id(&self) -> NodeId {
        self.id
    }

    pub fn shape(&self) -> Shape {
        self.handle.with_graph(|g| g.shape(self.id).clone())
    }

    pub fn dtype(&self) -> DType {
        self.shape().dtype()
    }

    pub fn rank(&self) -> usize {
        self.shape().rank()
    }

    /// Static dimensions as a `Vec<usize>` (panics on a dynamic axis).
    pub fn dims(&self) -> Vec<usize> {
        self.shape()
            .dims()
            .iter()
            .map(|d| match d {
                Dim::Static(n) => *n,
                Dim::Dynamic(_) => panic!("dims(): tensor has a dynamic axis"),
            })
            .collect()
    }

    /// Total number of elements (panics on a dynamic axis).
    pub fn numel(&self) -> usize {
        self.dims().iter().product()
    }

    /// True when every dimension is statically known.
    pub fn is_static_shape(&self) -> bool {
        self.shape().is_static()
    }

    /// True if `self` and `other` are backed by the **same** graph — i.e. one
    /// is a [`clone`](Clone) or a view (`slice`/`reshape`/`narrow`/…) of the
    /// other, sharing storage with zero copy. Tensors from independent
    /// constructors return `false` until they're combined in an op.
    pub fn shares_graph(&self, other: &Tensor) -> bool {
        self.handle.same(&other.handle)
    }

    /// Total bytes of constant (payload) data physically held in the backing
    /// graph. **Cloning and slicing do not increase this** — they share
    /// storage and only add zero-payload view nodes. Combining tensors from
    /// different graphs copies the right-hand operand's data in (once).
    pub fn storage_bytes(&self) -> usize {
        self.handle.with_graph(|g| {
            g.nodes()
                .iter()
                .filter_map(|n| match &n.op {
                    Op::Constant { data } => Some(data.len()),
                    _ => None,
                })
                .sum()
        })
    }

    /// View slice — lowers to [`narrow`](Self::narrow) (no copy at IR level).
    pub fn slice(&self, spec: SliceSpec) -> Self {
        spec.apply(self)
    }

    fn map_unary(&self, f: impl FnOnce(&mut rlx_ir::Graph, NodeId) -> NodeId) -> Self {
        let id = self.handle.with_graph(|g| f(g, self.id));
        Self::new(self.handle.clone(), id)
    }

    /// Pull `other`'s graph into this tensor's graph (no-op when they already
    /// share one), returning the operand's node id in *this* graph.
    pub(crate) fn adopt(&self, other: &Tensor) -> NodeId {
        self.handle.adopt(&other.handle, other.id)
    }

    fn map_binary(
        &self,
        rhs: BinaryRhs<'_>,
        f: impl FnOnce(&mut rlx_ir::Graph, NodeId, NodeId) -> NodeId,
    ) -> Self {
        let rhs_id = match rhs {
            BinaryRhs::Tensor(v) => self.adopt(v),
            BinaryRhs::Scalar(s) => self.handle.with_graph(|g| {
                let dtype = g.shape(self.id).dtype();
                promote_scalar(g, s, dtype)
            }),
        };
        let id = self.handle.with_graph(|g| f(g, self.id, rhs_id));
        Self::new(self.handle.clone(), id)
    }

    /// Apply a unary [`Activation`] with same-shape output.
    fn map_act(&self, act: Activation) -> Self {
        self.map_unary(|g, x| {
            let s = rlx_ir::shape::unary_shape(g.shape(x));
            g.activation(act, x, s)
        })
    }

    /// Elementwise [`BinaryOp`] with broadcast shape inference.
    fn map_binary_op(&self, rhs: BinaryRhs<'_>, op: BinaryOp) -> Self {
        self.map_binary(rhs, |g, a, b| {
            let s = rlx_ir::shape::binary_shape(g.shape(a), g.shape(b))
                .expect("binary shape inference");
            g.binary(op, a, b, s)
        })
    }

    /// Elementwise comparison ([`CmpOp`]) → `Bool` tensor.
    fn map_cmp(&self, rhs: BinaryRhs<'_>, op: CmpOp) -> Self {
        self.map_binary(rhs, |g, a, b| {
            let s = rlx_ir::shape::compare_shape(g.shape(a), g.shape(b))
                .expect("compare shape inference");
            g.add_node(Op::Compare(op), vec![a, b], s)
        })
    }

    /// Reduction ([`ReduceOp`]) over `axes`.
    fn map_reduce(&self, op: ReduceOp, axes: Vec<usize>, keep_dim: bool) -> Self {
        self.map_unary(|g, x| {
            let s = rlx_ir::shape::reduce_shape(g.shape(x), &axes, keep_dim)
                .expect("reduce shape inference");
            g.reduce(x, op, axes, keep_dim, s)
        })
    }

    /// Matrix multiply (`@` in Python / JAX).
    #[doc(alias = "@")]
    pub fn mm(&self, rhs: &Tensor) -> Self {
        self.matmul(rhs)
    }

    pub fn matmul(&self, rhs: &Tensor) -> Self {
        self.map_binary(BinaryRhs::Tensor(rhs), |g, a, b| g.mm(a, b))
    }

    /// Native low-precision GEMM (TN: `self [m,k] · rhs [n,k]ᵀ → [m,n]`).
    /// Both operands are dynamically quantized to the minifloat `fmt` +
    /// `layout`; `rhs` must be K-last (`[n, k]`). `fmt` may be any
    /// [`ScaledFormat`], including a parameterized `Custom` — e.g.
    /// `t.scaled_matmul(&w, ScaledFormat::custom(3, 0), ScaleLayout::mx())` runs
    /// the matmul in `f4e3m0`.
    pub fn scaled_matmul(&self, rhs: &Tensor, fmt: ScaledFormat, layout: ScaleLayout) -> Self {
        self.map_binary(BinaryRhs::Tensor(rhs), move |g, a, b| {
            g.scaled_matmul(a, b, fmt, layout)
        })
    }

    pub fn softmax(&self, axis: i32) -> Self {
        self.map_unary(|g, x| g.sm(x, axis))
    }

    pub fn reshape(&self, new_shape: impl Into<Vec<i64>>) -> Self {
        self.map_unary(|g, x| g.reshape_(x, new_shape.into()))
    }

    pub fn transpose(&self, perm: impl Into<Vec<usize>>) -> Self {
        self.map_unary(|g, x| g.transpose_(x, perm.into()))
    }

    pub fn t(&self) -> Self {
        let rank = self.rank();
        assert!(rank >= 2, "t() requires rank >= 2");
        let perm: Vec<usize> = (0..rank - 2).chain([rank - 1, rank - 2]).collect();
        self.transpose(perm)
    }

    pub fn cast(&self, to: DType) -> Self {
        self.map_unary(|g, x| g.cast(x, to))
    }

    pub fn narrow(&self, axis: usize, start: usize, len: usize) -> Self {
        self.map_unary(|g, x| g.narrow_(x, axis, start, len))
    }

    /// Broadcast (expand) to an explicit target shape — NumPy `broadcast_to`.
    /// A `1` axis stretches; sizes must otherwise match. Use `-1`-free static
    /// targets.
    pub fn broadcast_to(&self, shape: impl Into<Vec<i64>>) -> Self {
        let target = shape.into();
        self.map_unary(|g, x| {
            let s = rlx_ir::shape::expand_shape(g.shape(x), &target)
                .expect("broadcast_to shape inference");
            g.add_node(
                Op::Expand {
                    target_shape: target,
                },
                vec![x],
                s,
            )
        })
    }

    /// Flatten to a 1-D tensor (row-major).
    pub fn flatten(&self) -> Self {
        self.reshape(vec![-1_i64])
    }

    /// Insert a size-1 axis at `axis` (NumPy `expand_dims` / PyTorch
    /// `unsqueeze`).
    pub fn unsqueeze(&self, axis: usize) -> Self {
        let mut dims = self.dims_i64();
        assert!(axis <= dims.len(), "unsqueeze: axis {axis} out of range");
        dims.insert(axis, 1);
        self.reshape(dims)
    }

    /// Remove the size-1 axis at `axis` (panics if it is not size 1).
    pub fn squeeze(&self, axis: usize) -> Self {
        let dims = self.shape();
        let dims = dims.dims();
        assert!(
            matches!(dims.get(axis), Some(Dim::Static(1))),
            "squeeze: axis {axis} is not size 1"
        );
        let new: Vec<i64> = dims
            .iter()
            .enumerate()
            .filter(|(i, _)| *i != axis)
            .map(|(_, d)| dim_to_i64(d))
            .collect();
        self.reshape(new)
    }

    /// Remove every size-1 axis.
    pub fn squeeze_all(&self) -> Self {
        let new: Vec<i64> = self
            .shape()
            .dims()
            .iter()
            .filter(|d| !matches!(d, Dim::Static(1)))
            .map(dim_to_i64)
            .collect();
        self.reshape(new)
    }

    /// Current dims as `i64` (dynamic → `-1`), for reshape-style builders.
    fn dims_i64(&self) -> Vec<i64> {
        self.shape().dims().iter().map(dim_to_i64).collect()
    }

    /// Split `axis` into contiguous pieces of the given `sizes` (must sum to
    /// the axis length). NumPy `split` / PyTorch `split_with_sizes`.
    pub fn split(&self, axis: usize, sizes: &[usize]) -> Vec<Tensor> {
        let mut out = Vec::with_capacity(sizes.len());
        let mut start = 0;
        for &len in sizes {
            out.push(self.narrow(axis, start, len));
            start += len;
        }
        out
    }

    /// Split `axis` into `n` near-equal chunks (first `len % n` get one extra).
    pub fn chunk(&self, axis: usize, n: usize) -> Vec<Tensor> {
        assert!(n >= 1, "chunk: n must be >= 1");
        let total = self.dims()[axis];
        let base = total / n;
        let rem = total % n;
        let sizes: Vec<usize> = (0..n).map(|i| base + usize::from(i < rem)).collect();
        self.split(axis, &sizes)
    }

    /// Repeat the whole tensor `reps` times along `axis` (NumPy `tile` on one
    /// axis): output axis length becomes `len * reps`.
    pub fn tile(&self, axis: usize, reps: usize) -> Tensor {
        assert!(reps >= 1, "tile: reps must be >= 1");
        let copies = vec![self; reps];
        crate::array::cat(&copies, axis)
    }

    /// Reverse the order of elements along `axis` (NumPy `flip`).
    pub fn flip(&self, axis: usize) -> Tensor {
        let n = self.dims()[axis];
        let rev: Vec<i64> = (0..n as i64).rev().collect();
        self.index_select(axis, &Tensor::index_vec(rev))
    }

    /// Cyclically shift elements along `axis` by `shift` (NumPy `roll`).
    pub fn roll(&self, axis: usize, shift: usize) -> Tensor {
        let n = self.dims()[axis];
        let s = if n == 0 { 0 } else { shift % n };
        if s == 0 {
            return self.clone();
        }
        let tail = self.narrow(axis, n - s, s);
        let head = self.narrow(axis, 0, n - s);
        crate::array::cat(&[&tail, &head], axis)
    }

    /// Pad `axis` with `before`/`after` entries of `value` (constant pad).
    pub fn pad(&self, axis: usize, before: usize, after: usize, value: f32) -> Tensor {
        let mut parts: Vec<Tensor> = Vec::new();
        if before > 0 {
            let mut d = self.dims();
            d[axis] = before;
            parts.push(Tensor::full(&d, value));
        }
        parts.push(self.clone());
        if after > 0 {
            let mut d = self.dims();
            d[axis] = after;
            parts.push(Tensor::full(&d, value));
        }
        let refs: Vec<&Tensor> = parts.iter().collect();
        crate::array::cat(&refs, axis)
    }

    pub fn gather(&self, indices: &Tensor, axis: usize) -> Self {
        self.map_binary(BinaryRhs::Tensor(indices), |g, table, idx| {
            g.gather_(table, idx, axis)
        })
    }

    /// Gather along `axis` by an index tensor — NumPy `take` / PyTorch
    /// `index_select`. Alias for [`gather`](Self::gather).
    pub fn index_select(&self, axis: usize, indices: &Tensor) -> Self {
        self.gather(indices, axis)
    }

    /// Pick a single position along `axis`, **dropping** that axis (NumPy
    /// `a[…, i, …]`). Negative `index` counts from the end.
    pub fn select(&self, axis: usize, index: i64) -> Self {
        let dim = match self.shape().dims().get(axis) {
            Some(Dim::Static(n)) => Some(*n),
            _ => None,
        };
        let pos = crate::slice::resolve_index(index, dim, axis);
        self.narrow(axis, pos, 1).squeeze(axis)
    }

    pub fn stop_gradient(&self) -> Self {
        self.map_unary(|g, x| g.stop_gradient(x))
    }

    /// Index of the maximum along `axis` (f32-encoded indices). Drops the axis
    /// unless `keep_dim`.
    pub fn argmax(&self, axis: usize, keep_dim: bool) -> Self {
        self.map_unary(|g, x| {
            let s = rlx_ir::shape::reduce_shape(g.shape(x), &[axis], keep_dim)
                .expect("argmax shape inference");
            g.argmax(x, axis, keep_dim, s)
        })
    }

    /// Index of the minimum along `axis` (f32-encoded indices).
    pub fn argmin(&self, axis: usize, keep_dim: bool) -> Self {
        self.map_unary(|g, x| {
            let s = rlx_ir::shape::reduce_shape(g.shape(x), &[axis], keep_dim)
                .expect("argmin shape inference");
            g.argmin(x, axis, keep_dim, s)
        })
    }

    /// Matrix inverse of a square 2-D tensor. Composite over `DenseSolve`
    /// (`inv(A) = solve(A, I)`), so it needs a BLAS-backed eval build (LAPACK)
    /// to run, not the default splat backend.
    pub fn inv(&self) -> Self {
        let dims = self.dims();
        assert!(
            dims.len() == 2 && dims[0] == dims[1],
            "inv() requires a square 2-D matrix, got {dims:?}"
        );
        let n = dims[0];
        let eye_id = self.adopt(&Tensor::eye(n));
        let id = self.handle.with_graph(|g| {
            let s = Shape::new(&[n, n], DType::F32);
            g.dense_solve(self.id, eye_id, s)
        });
        Self::new(self.handle.clone(), id)
    }

    /// Solve the linear system `self @ x = b` for `x` (LAPACK via `DenseSolve`;
    /// needs a BLAS eval build). `x` takes `b`'s shape.
    pub fn solve(&self, b: &Tensor) -> Self {
        let s = b.shape();
        let b_id = self.adopt(b);
        let id = self.handle.with_graph(|g| g.dense_solve(self.id, b_id, s));
        Self::new(self.handle.clone(), id)
    }

    /// Cumulative sum along `axis`. `exclusive` shifts so `output[0]` = 0.
    pub fn cumsum(&self, axis: i32, exclusive: bool) -> Self {
        self.map_unary(|g, x| {
            let s = rlx_ir::shape::unary_shape(g.shape(x));
            g.cumsum(x, axis, exclusive, s)
        })
    }

    /// Population variance over `axes`: `mean((x - mean(x))²)`.
    pub fn var(&self, axes: impl Into<Vec<usize>>, keep_dim: bool) -> Self {
        let axes = axes.into();
        let mu = self.mean(axes.clone(), true);
        let centered = self - &mu;
        (&centered * &centered).mean(axes, keep_dim)
    }

    /// Population standard deviation over `axes`.
    pub fn std(&self, axes: impl Into<Vec<usize>>, keep_dim: bool) -> Self {
        self.var(axes, keep_dim).sqrt()
    }

    /// L2 (Frobenius) norm over `axes`: `sqrt(sum(x²))`.
    pub fn norm(&self, axes: impl Into<Vec<usize>>, keep_dim: bool) -> Self {
        (self * self).sum(axes, keep_dim).sqrt()
    }

    /// Numerically-stable `log(sum(exp(x)))` along `axis`
    /// (`max + log(sum(exp(x - max)))`).
    pub fn logsumexp(&self, axis: usize, keep_dim: bool) -> Self {
        let m = self.max([axis], true);
        let shifted = self - &m;
        let summed = shifted.exp().sum([axis], true);
        let lse = &summed.log() + &m;
        if keep_dim { lse } else { lse.squeeze(axis) }
    }

    /// 1-D FFT along the last axis (unnormalized; `ifft(fft(x)) = N·x`).
    /// Complex data uses the `[..., 2N]` block layout (real plane then imag
    /// plane) for F32/F64, or a `C64` last axis. Output matches input shape.
    pub fn fft(&self) -> Self {
        self.map_unary(|g, x| g.fft(x, false))
    }

    /// Inverse 1-D FFT along the last axis (unnormalized — pairs with [`Self::fft`]
    /// up to the `N` factor).
    pub fn ifft(&self) -> Self {
        self.map_unary(|g, x| g.fft(x, true))
    }

    /// 1-D FFT along an arbitrary `axis` (transpose → fft → transpose).
    pub fn fft_axis(&self, axis: usize, inverse: bool) -> Self {
        self.map_unary(|g, x| g.fft_axis(x, axis, inverse))
    }

    pub fn layer_norm(&self, gamma: &Tensor, beta: &Tensor, eps: f32) -> Self {
        let g_id = self.adopt(gamma);
        let b_id = self.adopt(beta);
        let id = self.handle.with_graph(|g| {
            let s = rlx_ir::shape::unary_shape(g.shape(self.id));
            g.layer_norm(self.id, g_id, b_id, -1, eps, s)
        });
        Self::new(self.handle.clone(), id)
    }

    pub fn rms_norm(&self, gamma: &Tensor, beta: &Tensor, eps: f32) -> Self {
        let g_id = self.adopt(gamma);
        let b_id = self.adopt(beta);
        let id = self
            .handle
            .with_graph(|g| g.rms_norm(self.id, g_id, b_id, eps));
        Self::new(self.handle.clone(), id)
    }

    /// Fused softmax + cross-entropy against a dense (soft / one-hot) target
    /// distribution — the standard classifier / language-model loss, in one
    /// differentiable op (native CPU & Metal kernels; lowered to primitives on
    /// other backends). `self` are logits `[N, C]`, `targets` a probability
    /// distribution `[N, C]` (one-hot, or label-smoothed). Returns the per-row
    /// loss `[N]`: `loss[n] = logsumexp(logits[n]) − Σ_c targets[n,c]·logits[n,c]`.
    /// Reduce it to a scalar training loss with [`mean_all`](Self::mean_all) —
    /// `mean(cross_entropy(logits, targets))` in the `rlx!` DSL. Gradients flow
    /// to both `self` and `targets`.
    pub fn cross_entropy(&self, targets: &Tensor) -> Self {
        let t_id = self.adopt(targets);
        let id = self
            .handle
            .with_graph(|g| g.softmax_cross_entropy(self.id, t_id));
        Self::new(self.handle.clone(), id)
    }

    /// Fused softmax + cross-entropy with integer class labels — PyTorch
    /// `cross_entropy(logits, labels, reduction="none")`. `self` are logits
    /// `[N, C]`, `labels` the f32-encoded class ids `[N]`. Returns the per-row
    /// loss `[N]`; reduce with [`mean_all`](Self::mean_all) for a scalar loss.
    pub fn softmax_cross_entropy_with_logits(&self, labels: &Tensor) -> Self {
        let l_id = self.adopt(labels);
        let id = self
            .handle
            .with_graph(|g| g.softmax_cross_entropy_with_logits(self.id, l_id));
        Self::new(self.handle.clone(), id)
    }

    /// Mean over *every* axis → a scalar (rank-0). The idiomatic loss reduction
    /// — `mean(cross_entropy(logits, targets))` in the `rlx!` DSL, or
    /// `loss.mean_all()` in Rust. A rank-0 input is returned unchanged.
    pub fn mean_all(&self) -> Self {
        let r = self.rank();
        if r == 0 {
            return self.clone();
        }
        self.mean((0..r).collect::<Vec<usize>>(), false)
    }

    /// Sum over *every* axis → a scalar (rank-0). A rank-0 input is returned
    /// unchanged.
    pub fn sum_all(&self) -> Self {
        let r = self.rank();
        if r == 0 {
            return self.clone();
        }
        self.sum((0..r).collect::<Vec<usize>>(), false)
    }

    pub fn conv2d(
        &self,
        weight: &Tensor,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        groups: usize,
    ) -> Self {
        let w_id = self.adopt(weight);
        let id = self.handle.with_graph(|g| {
            g.conv2d(
                self.id,
                w_id,
                kernel_size,
                stride,
                padding,
                dilation,
                groups,
            )
        });
        Self::new(self.handle.clone(), id)
    }

    pub fn attention(
        &self,
        k: &Tensor,
        v: &Tensor,
        num_heads: usize,
        head_dim: usize,
        mask_kind: MaskKind,
    ) -> Self {
        let k_id = self.adopt(k);
        let v_id = self.adopt(v);
        let id = self.handle.with_graph(|g| {
            let s = rlx_ir::shape::attention_shape(g.shape(self.id));
            g.attention_kind(self.id, k_id, v_id, num_heads, head_dim, mask_kind, s)
        });
        Self::new(self.handle.clone(), id)
    }

    /// On-chip codebook weight-synthesis matmul (`Op::SynthMatMul`): `y = x·Wᵀ`
    /// where the weight `[n, k]` is reconstructed inside the matmul from `indices`
    /// `[n, k/entry_dim]` (u8, fixed quantization structure) + `codebook`
    /// `[num_entries, entry_dim]` (f32, trained). `x` is `[m, k]`, output `[m, n]`.
    /// The codebook is a tiny trainable param; the full weight is never stored —
    /// "functions not data". Backs `rlx-tiny`.
    pub fn synth_matmul(
        &self,
        indices: &Tensor,
        codebook: &Tensor,
        entry_dim: u32,
        num_entries: u32,
    ) -> Self {
        let idx_id = self.adopt(indices);
        let cb_id = self.adopt(codebook);
        let id = self.handle.with_graph(|g| {
            let m = g.shape(self.id).dim(0).unwrap_static();
            let n = g.shape(idx_id).dim(0).unwrap_static();
            let out = rlx_ir::Shape::new(&[m, n], rlx_ir::DType::F32);
            g.synth_matmul(
                self.id,
                idx_id,
                cb_id,
                rlx_ir::SynthKind::Codebook {
                    entry_dim,
                    num_entries,
                },
                out,
            )
        });
        Self::new(self.handle.clone(), id)
    }

    /// Reconstruct the dense weight `W [k,n]` from `self` = indices `[n, k/entry_dim]`
    /// + `codebook [num_entries, entry_dim]` (the same gather the SynthMatMul
    /// decompose uses, WITHOUT the matmul). Lets a model fold multi-stage VQ + LoRA
    /// into ONE effective weight (`W_eff = Σ reconstruct + A·Bᵀ`) so `x·W_eff` is a
    /// single GEMM — dense-speed forward AND a 2-GEMM backward (the multi-stage /
    /// LoRA gradients collapse to codebook scatters + rank-r products with no batch
    /// dim). For training/prefill (large m); decode keeps `synth_matmul` in-kernel.
    pub fn synth_reconstruct(&self, codebook: &Tensor, entry_dim: u32) -> Self {
        let cb_id = self.adopt(codebook);
        let id = self.handle.with_graph(|g| {
            use rlx_ir::{DType, Op, Shape};
            // DECOMPOSED reconstruction (Cast→Reshape→Gather→Reshape→Transpose →
            // W[k,n]), NOT the native `Op::SynthReconstruct` fused kernel. MEASURED
            // (both `[k,n]` and `[n,k]` native layouts): the fused kernel wins the
            // FORWARD (~35 vs 47ms — one dispatch vs cast+gather+reshape) but LOSES
            // the backward by more (~+20ms), because the opaque op hides structure
            // from the backward CSE + transpose-simplification that dedups the
            // per-stage `grad_W` transposes and makes `dx=dy·w_bt` free. Net-worse
            // for TRAINING → keep the transparent primitives here. The native op +
            // kernel + `LowerSynthReconstruct` oracle stay in-tree (dormant) for a
            // future INFERENCE-only path (forward win, no backward to regress).
            let d = entry_dim as usize;
            let idx_shape = g.shape(self.id);
            let n = idx_shape.dim(0).unwrap_static();
            let kb = idx_shape.dim(1).unwrap_static();
            let (k, p) = (kb * d, n * kb);
            let idx_i64 = g.add_node(
                Op::Cast { to: DType::I64 },
                vec![self.id],
                Shape::new(&[n, kb], DType::I64),
            );
            let idx_flat = g.add_node(
                Op::Reshape {
                    new_shape: vec![p as i64],
                },
                vec![idx_i64],
                Shape::new(&[p], DType::I64),
            );
            let rows = g.add_node(
                Op::Gather { axis: 0 },
                vec![cb_id, idx_flat],
                Shape::new(&[p, d], DType::F32),
            );
            let w_bt = g.add_node(
                Op::Reshape {
                    new_shape: vec![n as i64, k as i64],
                },
                vec![rows],
                Shape::new(&[n, k], DType::F32),
            );
            g.add_node(
                Op::Transpose { perm: vec![1, 0] },
                vec![w_bt],
                Shape::new(&[k, n], DType::F32),
            )
        });
        Self::new(self.handle.clone(), id)
    }

    /// KAN learnable per-channel spline activation (`Op::SplineActivation`): each
    /// channel `c` of the last axis gets its own univariate function, a Gaussian-RBF
    /// expansion `Σ_g coeff[c,g]·exp(-((x−center_g)·inv_h)²)`. `coeff` is
    /// `[C, num_basis]` (trained); output matches `self`. A learnable activation
    /// replacing a fixed GELU/ReLU. Backs `rlx-tiny`.
    pub fn spline_activation(
        &self,
        coeff: &Tensor,
        num_basis: u32,
        grid_min: f32,
        grid_max: f32,
    ) -> Self {
        let c_id = self.adopt(coeff);
        let id = self
            .handle
            .with_graph(|g| g.spline_activation(self.id, c_id, num_basis, grid_min, grid_max));
        Self::new(self.handle.clone(), id)
    }

    /// HuggingFace-style linear `x · Wᵀ + b` where `w` is stored `[out, in]`,
    /// as a single fused matmul+bias node (`Op::FusedMatMulBiasAct`). The
    /// weight transpose is emitted in-graph so callers bind the raw `[out, in]`
    /// checkpoint tensor directly (no host pre-transpose). Backs the `rlx!`
    /// `linear(x, w, b)` / `x @t w` sugar.
    pub fn linear(&self, w: &Tensor, b: &Tensor) -> Self {
        self.linear_act(w, b, None)
    }

    /// [`linear`](Self::linear) with a fused elementwise activation applied to
    /// the result — `act(x · Wᵀ + b)` in one op. Backs `gelu(linear(..))` etc.
    pub fn linear_act(&self, w: &Tensor, b: &Tensor, act: Option<Activation>) -> Self {
        let w_id = self.adopt(w);
        let b_id = self.adopt(b);
        let id = self.handle.with_graph(|g| {
            let rank = g.shape(w_id).rank();
            let wt = g.transpose_(w_id, vec![rank - 1, rank - 2]);
            let s = rlx_ir::shape::matmul_shape(g.shape(self.id), g.shape(wt))
                .expect("linear shape inference");
            g.fused_matmul_bias_act(self.id, wt, b_id, act, s)
        });
        Self::new(self.handle.clone(), id)
    }

    /// Matrix multiply against a transposed right operand: `self · rhsᵀ`
    /// (`x @t w` in the DSL). Avoids a separate `.t()` at call sites; the
    /// transpose is emitted in-graph.
    pub fn matmul_t(&self, rhs: &Tensor) -> Self {
        self.matmul(&rhs.t())
    }

    /// Scaled dot-product attention with an explicit `[batch, key_len]` padding
    /// mask (`1.0` = valid, `<0.5` = ignored) — the 4-input `MaskKind::Custom`
    /// form that [`attention`](Self::attention) (kernel-synthesized masks) can't
    /// express. Backs the `rlx!` `attention(q, k, v, mask, heads, head_dim)`
    /// sugar.
    pub fn attention_masked(
        &self,
        k: &Tensor,
        v: &Tensor,
        mask: &Tensor,
        num_heads: usize,
        head_dim: usize,
    ) -> Self {
        let k_id = self.adopt(k);
        let v_id = self.adopt(v);
        let m_id = self.adopt(mask);
        let id = self.handle.with_graph(|g| {
            let s = rlx_ir::shape::attention_shape(g.shape(self.id));
            g.attention(self.id, k_id, v_id, m_id, num_heads, head_dim, s)
        });
        Self::new(self.handle.clone(), id)
    }

    pub fn rope(&self, cos: &Tensor, sin: &Tensor, head_dim: usize) -> Self {
        let cos_id = self.adopt(cos);
        let sin_id = self.adopt(sin);
        let id = self
            .handle
            .with_graph(|g| g.rope(self.id, cos_id, sin_id, head_dim));
        Self::new(self.handle.clone(), id)
    }

    /// Select `on_true` where `self` (a mask) is true, else `on_false`. Each
    /// branch is a `&Tensor` or a scalar (`mask.where_(&x, 0.0)`), so the
    /// `rlx!` DSL's `select(mask, x, 0.0)` lowers here directly. A scalar
    /// branch is promoted to the other branch's dtype (else `F32`).
    pub fn where_<'t, 'f>(
        &self,
        on_true: impl Into<BinaryRhs<'t>>,
        on_false: impl Into<BinaryRhs<'f>>,
    ) -> Self {
        let t = on_true.into();
        let f = on_false.into();
        let hint = if let BinaryRhs::Tensor(x) = &t {
            x.dtype()
        } else if let BinaryRhs::Tensor(x) = &f {
            x.dtype()
        } else {
            DType::F32
        };
        let t_id = self.where_branch(t, hint);
        let f_id = self.where_branch(f, hint);
        let id = self.handle.with_graph(|g| {
            let s = rlx_ir::shape::binary_shape(g.shape(t_id), g.shape(f_id))
                .expect("where shape inference");
            g.add_node(Op::Where, vec![self.id, t_id, f_id], s)
        });
        Self::new(self.handle.clone(), id)
    }

    /// Resolve a `where_` branch to a node id in this graph — adopting a tensor
    /// or promoting a scalar to `hint`'s dtype.
    fn where_branch(&self, rhs: BinaryRhs<'_>, hint: DType) -> NodeId {
        match rhs {
            BinaryRhs::Tensor(t) => self.adopt(t),
            BinaryRhs::Scalar(s) => self.handle.with_graph(|g| promote_scalar(g, s, hint)),
        }
    }

    /// Replace entries of `self` where `mask` is true with `value` (NumPy
    /// `masked_fill` / the attention-mask primitive: `scores.masked_fill(causal,
    /// f32::NEG_INFINITY)` then `softmax`). Requires a static shape.
    pub fn masked_fill(&self, mask: &Tensor, value: f64) -> Self {
        let dims: Vec<usize> = self
            .shape()
            .dims()
            .iter()
            .map(|d| match d {
                Dim::Static(n) => *n,
                Dim::Dynamic(_) => panic!("masked_fill requires a static shape"),
            })
            .collect();
        let fill = Tensor::full(dims, value as f32);
        mask.where_(&fill, self)
    }

    /// Elementwise maximum (accepts a `&Tensor` or a scalar).
    pub fn maximum<'r>(&self, rhs: impl Into<BinaryRhs<'r>>) -> Self {
        self.map_binary_op(rhs.into(), BinaryOp::Max)
    }

    /// Elementwise minimum (accepts a `&Tensor` or a scalar).
    pub fn minimum<'r>(&self, rhs: impl Into<BinaryRhs<'r>>) -> Self {
        self.map_binary_op(rhs.into(), BinaryOp::Min)
    }

    /// Elementwise `atan2(self, rhs)` — quadrant-aware arctangent.
    pub fn atan2<'r>(&self, rhs: impl Into<BinaryRhs<'r>>) -> Self {
        self.map_binary_op(rhs.into(), BinaryOp::Atan2)
    }

    /// Elementwise `self % rhs` (C `fmod`; accepts a `&Tensor` or a scalar).
    pub fn rem<'r>(&self, rhs: impl Into<BinaryRhs<'r>>) -> Self {
        self.map_binary_op(rhs.into(), BinaryOp::Mod)
    }

    /// Clamp values to `[min, max]`.
    pub fn clamp(&self, min: f64, max: f64) -> Self {
        self.clamp_min(min).clamp_max(max)
    }

    /// Lower-bound values at `min`.
    pub fn clamp_min(&self, min: f64) -> Self {
        self.map_binary_op(BinaryRhs::Scalar(min.into()), BinaryOp::Max)
    }

    /// Upper-bound values at `max`.
    pub fn clamp_max(&self, max: f64) -> Self {
        self.map_binary_op(BinaryRhs::Scalar(max.into()), BinaryOp::Min)
    }

    /// Elementwise power `self ^ rhs` (accepts a `&Tensor` or a scalar).
    pub fn pow<'r>(&self, rhs: impl Into<BinaryRhs<'r>>) -> Self {
        self.map_binary_op(rhs.into(), BinaryOp::Pow)
    }

    pub fn pow_scalar(&self, exp: f64) -> Self {
        self.map_binary_op(BinaryRhs::Scalar(exp.into()), BinaryOp::Pow)
    }
}

unary_ops! {
    /// Rectified linear unit `max(x, 0)`.
    relu => Relu,
    /// Gaussian error linear unit (exact).
    gelu => Gelu,
    /// GELU, tanh approximation.
    gelu_approx => GeluApprox,
    /// Sigmoid-weighted linear unit `x · σ(x)` (a.k.a. SwiGLU gate).
    silu => Silu,
    /// Hyperbolic tangent.
    tanh => Tanh,
    /// Natural exponential `e^x`.
    exp => Exp,
    /// Square root.
    sqrt => Sqrt,
    /// Logistic sigmoid `1 / (1 + e^-x)`.
    sigmoid => Sigmoid,
    /// Natural logarithm.
    log => Log,
    /// Reciprocal square root `1 / sqrt(x)`.
    rsqrt => Rsqrt,
    /// Absolute value.
    abs => Abs,
    /// Sine.
    sin => Sin,
    /// Cosine.
    cos => Cos,
    /// Tangent.
    tan => Tan,
    /// Arctangent.
    atan => Atan,
    /// Round to nearest (half-to-even); straight-through gradient.
    round => Round,
    /// Reciprocal `1 / x`.
    recip => Recip,
    /// `floor(x)` (piecewise-constant gradient).
    floor => Floor,
    /// `ceil(x)`.
    ceil => Ceil,
    /// `sign(x)` = -1/0/+1.
    sign => Sign,
    /// Softplus `ln(1 + eˣ)`.
    softplus => Softplus,
    /// ELU (α=1): `x if x>0 else eˣ-1`.
    elu => Elu,
    /// Gauss error function `erf(x)`.
    erf => Erf,
    /// HardSwish `x·clamp(x+3,0,6)/6`.
    hard_swish => HardSwish,
    /// HardSigmoid `clamp(x/6+0.5,0,1)`.
    hard_sigmoid => HardSigmoid,
    /// Mish `x·tanh(softplus(x))`.
    mish => Mish,
    /// Softsign `x/(1+|x|)`.
    softsign => Softsign,
    /// LogSigmoid `log(σ(x))`.
    log_sigmoid => LogSigmoid,
}

cmp_ops! {
    /// Elementwise `==`.
    eq => Eq,
    /// Elementwise `!=`.
    ne => Ne,
    /// Elementwise `<`.
    lt => Lt,
    /// Elementwise `<=`.
    le => Le,
    /// Elementwise `>`.
    gt => Gt,
    /// Elementwise `>=`.
    ge => Ge,
}

reduce_ops! {
    /// Sum over `axes`.
    sum => Sum,
    /// Mean over `axes`.
    mean => Mean,
    /// Maximum over `axes`.
    max => Max,
    /// Minimum over `axes`.
    min => Min,
    /// Product over `axes`.
    prod => Prod,
}

/// Generate typed little-endian readback methods (`name => (T, DType, width)`):
/// realize, assert the tensor dtype, then decode fixed-width POD elements.
/// `f32` (the default) and `bool` (multi-encoding) are hand-rolled separately.
#[cfg(feature = "eval")]
macro_rules! pod_readers {
    ($($(#[$m:meta])* $name:ident => ($ty:ty, $dt:ident, $width:literal)),+ $(,)?) => {
        $(
            $(#[$m])*
            pub fn $name(&self) -> Vec<$ty> {
                let (bytes, dt) = self.eval_typed();
                assert_eq!(
                    dt, DType::$dt,
                    concat!(stringify!($name), ": tensor dtype is {:?}"), dt,
                );
                bytes
                    .chunks_exact($width)
                    .map(|c| <$ty>::from_le_bytes(c.try_into().unwrap()))
                    .collect()
            }
        )+
    };
}

/// Materialization — compile this tensor's graph through `rlx_runtime` and
/// read the result back to host memory. The graph stays lazy until one of
/// these is called, so the compiler fuses and memory-plans the whole
/// expression first. Available with the `eval` feature.
#[cfg(feature = "eval")]
impl Tensor {
    /// Compile + run and copy the output to a `Vec<f32>`. The device is chosen
    /// automatically — the fastest backend compiled into this build that can
    /// run the graph (CPU when nothing faster is available).
    pub fn to_vec(&self) -> Vec<f32> {
        self.to_vec_on(self.auto_device())
    }

    /// Compile + run on an explicit device.
    pub fn to_vec_on(&self, device: rlx_runtime::Device) -> Vec<f32> {
        // A raw constant's realized value is its own host data — device-independent,
        // so return it directly and skip backend compilation. This is both an
        // optimization (no GPU round-trip for a constant) and a robustness fix: the
        // MLX backend crashes when asked to compile a constant-only trace.
        if let Some((bytes, DType::F32)) = self.constant_payload() {
            return bytes
                .chunks_exact(4)
                .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
                .collect();
        }
        let compiled = self.compiled_on(device);
        let mut outputs = compiled.borrow_mut().run(&[]);
        outputs.pop().expect("eval produced no output")
    }

    /// If this tensor's output node is a raw `Op::Constant` that physically holds its
    /// full contiguous payload, returns `(bytes, dtype)` — its realized value with no
    /// computation to run. (Views/slices of a constant are *not* raw constants, so
    /// they fall through to normal evaluation.)
    fn constant_payload(&self) -> Option<(Vec<u8>, DType)> {
        self.handle.with_graph(|g| {
            let node = g.node(self.id);
            if let Op::Constant { data } = &node.op {
                if node.shape.size_bytes() == Some(data.len()) {
                    return Some((data.clone(), node.shape.dtype()));
                }
            }
            None
        })
    }

    /// Fastest available backend that can run this tensor's graph (no clone).
    fn auto_device(&self) -> rlx_runtime::Device {
        self.handle
            .with_graph(|g| rlx_runtime::fastest_device_for(g))
    }

    /// Compiled graph for this tensor's output — borrows the shared graph for
    /// the cache fingerprint, cloning only on a miss (zero-copy on hits).
    fn compiled_on(
        &self,
        device: rlx_runtime::Device,
    ) -> std::rc::Rc<std::cell::RefCell<rlx_runtime::CompiledGraph>> {
        self.handle
            .with_graph(|g| crate::cache::compiled_output(g, self.id, device))
    }

    /// Force evaluation — alias for [`to_vec`](Self::to_vec). Building a
    /// `Tensor` allocates **nothing** and runs **nothing**; it only records IR.
    /// `realize` is the single point where the accumulated graph is compiled
    /// (fused + memory-planned) and executed. Use it to make the lazy model
    /// explicit at call sites.
    pub fn realize(&self) -> Vec<f32> {
        self.to_vec()
    }

    /// [`realize`](Self::realize) on an explicit device.
    pub fn realize_on(&self, device: rlx_runtime::Device) -> Vec<f32> {
        self.to_vec_on(device)
    }

    /// Realize a scalar (rank-0 or single-element) tensor to one `f32`.
    pub fn item(&self) -> f32 {
        let v = self.to_vec();
        assert_eq!(v.len(), 1, "item(): tensor has {} elements, not 1", v.len());
        v[0]
    }

    pod_readers! {
        /// Realize and read back as `f64` (tensor dtype must be `F64`).
        to_vec_f64 => (f64, F64, 8),
        /// Realize and read back as `i64` (tensor dtype must be `I64`).
        to_vec_i64 => (i64, I64, 8),
        /// Realize and read back as `i32` (tensor dtype must be `I32`).
        to_vec_i32 => (i32, I32, 4),
    }

    /// Realize and read back a `Bool` tensor. Robust to the backend storing
    /// bools 1-byte or f32-encoded (an element is `true` iff any of its bytes
    /// is non-zero).
    pub fn to_vec_bool(&self) -> Vec<bool> {
        let (bytes, dt) = self.eval_typed();
        assert_eq!(dt, DType::Bool, "to_vec_bool: tensor dtype is {dt:?}");
        let n = self.numel();
        assert!(n > 0 && bytes.len() % n == 0, "to_vec_bool: ragged output");
        let w = bytes.len() / n;
        bytes
            .chunks_exact(w)
            .map(|c| c.iter().any(|&b| b != 0))
            .collect()
    }

    /// Compile + run, returning the output's raw little-endian bytes + dtype.
    /// Device is auto-selected; backs the typed `to_vec_*` readers.
    fn eval_typed(&self) -> (Vec<u8>, DType) {
        // Raw constant → its bytes are the result (see `to_vec_on`).
        if let Some(ct) = self.constant_payload() {
            return ct;
        }
        let compiled = self.compiled_on(self.auto_device());
        let mut outputs = compiled.borrow_mut().run_typed(&[]);
        outputs.pop().expect("eval produced no output")
    }

    /// Realize and pretty-print with shape + dtype, NumPy/ndarray-style nesting.
    /// `println!("{}", t)` does the same (with the `eval` feature).
    pub fn show(&self) -> String {
        let dims = self.dims();
        let data = self.to_vec();
        format!(
            "Tensor{dims:?} {:?}\n{}",
            self.dtype(),
            fmt_nested(&data, &dims)
        )
    }

    /// Select a device fluently: `t.on(Device::Metal).to_vec()`.
    pub fn on(&self, device: rlx_runtime::Device) -> Materialize<'_> {
        Materialize {
            tensor: self,
            device,
        }
    }
}

/// Fluent device selector returned by [`Tensor::on`].
#[cfg(feature = "eval")]
pub struct Materialize<'a> {
    tensor: &'a Tensor,
    device: rlx_runtime::Device,
}

#[cfg(feature = "eval")]
impl Materialize<'_> {
    /// Compile + run on the selected device and copy the output to host.
    pub fn to_vec(self) -> Vec<f32> {
        self.tensor.to_vec_on(self.device)
    }
}

/// Reverse-mode autodiff. Available with the `grad` feature; pair with `eval`
/// to read the resulting gradient values.
#[cfg(feature = "autodiff")]
impl Tensor {
    /// Gradients of this tensor (treated as the loss) w.r.t. each `wrt`,
    /// returned as symbolic [`Tensor`]s in the same order as `wrt`.
    ///
    /// The seed `∂self/∂self = 1` is baked in as a constant (ones, matching
    /// `self`'s shape — so a non-scalar `self` differentiates `sum(self)`),
    /// making each returned gradient a self-contained graph you can
    /// `to_vec()` or compose further.
    ///
    /// ```ignore
    /// let loss = (&a * &b).sum([0], false);
    /// let g = loss.grad(&[&a, &b]);   // [∂/∂a, ∂/∂b]
    /// assert_eq!(g[0].to_vec(), b.to_vec());
    /// ```
    pub fn grad(&self, wrt: &[&Tensor]) -> Vec<Tensor> {
        // Resolve each `wrt` to its node id *within this loss's graph*.
        // Memoized adoption guarantees these are the exact nodes the loss was
        // built from, so gradient actually flows (not to a stale duplicate).
        let wrt_ids: Vec<NodeId> = wrt.iter().map(|t| self.adopt(t)).collect();
        let loss_id = self.id;
        let bwd = self.handle.with_graph(|g| {
            let mut forward = g.clone();
            forward.set_outputs(vec![loss_id]);
            let mut bwd = rlx_autodiff::grad(&forward, &wrt_ids);
            bake_unit_seed(&mut bwd);
            bwd
        });
        let handle = GraphHandle::new(bwd);
        let outputs = handle.with_graph(|g| g.outputs.clone());
        outputs
            .into_iter()
            .map(|id| Tensor::new(handle.clone(), id))
            .collect()
    }
}

/// Replace `rlx_autodiff`'s `d_output` seed input with a constant of ones, so
/// the backward graph runs with no external inputs.
#[cfg(feature = "autodiff")]
pub(crate) fn bake_unit_seed(bwd: &mut rlx_ir::Graph) {
    let seed = bwd.nodes().iter().find_map(|n| match &n.op {
        Op::Input { name } if name == "d_output" => Some(n.id),
        _ => None,
    });
    if let Some(id) = seed {
        let n = bwd.node(id).shape.num_elements().unwrap_or(1);
        let data: Vec<u8> = (0..n).flat_map(|_| 1.0f32.to_le_bytes()).collect();
        let node = bwd.node_mut(id);
        node.op = Op::Constant { data };
        node.inputs = Vec::new();
    }
}

impl From<Tensor> for NodeId {
    fn from(v: Tensor) -> Self {
        v.id
    }
}

impl From<&Tensor> for NodeId {
    fn from(v: &Tensor) -> Self {
        v.id
    }
}

pub(crate) enum BinaryRhs<'a> {
    Tensor(&'a Tensor),
    Scalar(Scalar),
}

impl<'a> BinaryRhs<'a> {
    fn scalar(s: Scalar) -> Self {
        Self::Scalar(s)
    }
}

impl<'a> From<&'a Tensor> for BinaryRhs<'a> {
    fn from(v: &'a Tensor) -> Self {
        Self::Tensor(v)
    }
}

macro_rules! impl_scalar_rhs {
    ($($ty:ty),+ $(,)?) => {
        $(
            impl From<$ty> for BinaryRhs<'_> {
                fn from(v: $ty) -> Self {
                    Self::Scalar(v.into())
                }
            }
        )+
    };
}

impl_scalar_rhs!(bool, i32, i64, f32, f64);

macro_rules! impl_tensor_binop {
    ($trait:ident, $method:ident, $graph:ident) => {
        impl $trait<&Tensor> for &Tensor {
            type Output = Tensor;

            fn $method(self, rhs: &Tensor) -> Tensor {
                self.map_binary(BinaryRhs::Tensor(rhs), |g, a, b| g.$graph(a, b))
            }
        }

        impl $trait<Tensor> for &Tensor {
            type Output = Tensor;

            fn $method(self, rhs: Tensor) -> Tensor {
                self.$method(&rhs)
            }
        }

        impl $trait<f64> for &Tensor {
            type Output = Tensor;

            fn $method(self, rhs: f64) -> Tensor {
                self.map_binary(BinaryRhs::scalar(rhs.into()), |g, a, b| g.$graph(a, b))
            }
        }

        impl $trait<f32> for &Tensor {
            type Output = Tensor;

            fn $method(self, rhs: f32) -> Tensor {
                self.$method(rhs as f64)
            }
        }

        impl $trait<i32> for &Tensor {
            type Output = Tensor;

            fn $method(self, rhs: i32) -> Tensor {
                self.map_binary(BinaryRhs::scalar(rhs.into()), |g, a, b| g.$graph(a, b))
            }
        }

        impl $trait<&Tensor> for Tensor {
            type Output = Tensor;

            fn $method(self, rhs: &Tensor) -> Tensor {
                (&self).$method(rhs)
            }
        }

        impl $trait<Tensor> for Tensor {
            type Output = Tensor;

            fn $method(self, rhs: Tensor) -> Tensor {
                (&self).$method(&rhs)
            }
        }

        impl $trait<f64> for Tensor {
            type Output = Tensor;

            fn $method(self, rhs: f64) -> Tensor {
                (&self).$method(rhs)
            }
        }

        impl $trait<f32> for Tensor {
            type Output = Tensor;

            fn $method(self, rhs: f32) -> Tensor {
                (&self).$method(rhs)
            }
        }

        impl $trait<i32> for Tensor {
            type Output = Tensor;

            fn $method(self, rhs: i32) -> Tensor {
                (&self).$method(rhs)
            }
        }
    };
}

impl_tensor_binop!(Add, add, add);
impl_tensor_binop!(Sub, sub, sub);
impl_tensor_binop!(Mul, mul, mul);
impl_tensor_binop!(Div, div, div);

impl Neg for &Tensor {
    type Output = Tensor;

    fn neg(self) -> Tensor {
        self.map_unary(|g, x| g.neg(x))
    }
}

macro_rules! impl_scalar_left {
    ($trait:ident, $method:ident, $graph:ident) => {
        impl $trait<&Tensor> for f64 {
            type Output = Tensor;

            fn $method(self, rhs: &Tensor) -> Tensor {
                let id = rhs.handle.with_graph(|g| {
                    let lhs_id = promote_scalar(g, self.into(), g.shape(rhs.id).dtype());
                    g.$graph(lhs_id, rhs.id)
                });
                Tensor::new(rhs.handle.clone(), id)
            }
        }

        impl $trait<&Tensor> for f32 {
            type Output = Tensor;

            fn $method(self, rhs: &Tensor) -> Tensor {
                (self as f64).$method(rhs)
            }
        }

        impl $trait<&Tensor> for i32 {
            type Output = Tensor;

            fn $method(self, rhs: &Tensor) -> Tensor {
                (self as i64 as f64).$method(rhs)
            }
        }
    };
}

impl_scalar_left!(Add, add, add);
impl_scalar_left!(Sub, sub, sub);
impl_scalar_left!(Mul, mul, mul);
impl_scalar_left!(Div, div, div);

impl Neg for Tensor {
    type Output = Tensor;

    fn neg(self) -> Tensor {
        (&self).neg()
    }
}