oxmera_tensor/

ops.rs

//! The differentiable operation layer: every method dispatches to the
//! device backend for the forward pass and, when recording is on and an
//! input is tracked, attaches the exact vector-Jacobian product to the
//! output's tape node.

use oxmera_core::{DType, Device, Error, Result, Shape};

use crate::autograd::{GradFn, is_recording};
use crate::backend::{Backend, BinaryOp, ReduceOp, UnaryOp, backend_for};
use crate::tensor::{Tensor, ViewKind};

use std::sync::Arc;

fn same_device(a: &Tensor, b: &Tensor, op: &'static str) -> Result<Device> {
    if a.device() != b.device() {
        return Err(Error::DeviceMismatch {
            lhs: a.device(),
            rhs: b.device(),
            op,
        });
    }
    Ok(a.device())
}

fn record(
    out: Tensor,
    inputs: Vec<Tensor>,
    vjp: impl Fn(&Tensor) -> Result<Vec<Option<Tensor>>> + Send + Sync + 'static,
) -> Tensor {
    if is_recording() && inputs.iter().any(Tensor::is_tracked) {
        out.with_grad_fn(GradFn {
            inputs,
            vjp: Box::new(vjp),
        })
    } else {
        out
    }
}

/// Sum `grad` down to `shape` (undo broadcasting): reduce the leading
/// extra axes and every axis the target holds as 1, then reshape.
pub(crate) fn reduce_to_shape(grad: &Tensor, shape: &Shape) -> Result<Tensor> {
    if grad.shape() == shape {
        return Ok(grad.clone());
    }
    let gdims = grad.dims().to_vec();
    let tdims = shape.dims();
    let lead = gdims.len() - tdims.len();
    let mut axes: Vec<usize> = (0..lead).collect();
    for (i, &td) in tdims.iter().enumerate() {
        if td == 1 && gdims[lead + i] != 1 {
            axes.push(lead + i);
        }
    }
    let reduced = if axes.is_empty() {
        grad.clone()
    } else {
        grad.sum_keepdim(&axes, true)?
    };
    reduced.reshape(shape.clone())
}

/// Attach the view VJP to a freshly built view (called from `tensor.rs`).
pub(crate) fn record_view(input: &Tensor, out: Tensor, kind: ViewKind) -> Tensor {
    let in_shape = input.shape().clone();
    record(out, vec![input.clone()], move |g| {
        let gi = match &kind {
            ViewKind::Reshape | ViewKind::Contiguous => g.reshape(in_shape.clone())?,
            ViewKind::Permute(perm) => {
                let mut inverse = vec![0usize; perm.len()];
                for (i, &p) in perm.iter().enumerate() {
                    inverse[p] = i;
                }
                g.permute(&inverse)?
            }
            ViewKind::Narrow { dim, start, len } => {
                let indices: Vec<i64> = (*start..start + len).map(|i| i as i64).collect();
                let indices = Tensor::from_vec_i64(indices, Shape::from([*len]))?;
                Tensor::zeros(in_shape.clone())
                    .to_dtype(g.dtype())?
                    .to_device(g.device())?
                    .index_add(*dim, &indices, g)?
            }
            ViewKind::Broadcast => reduce_to_shape(g, &in_shape)?,
        };
        Ok(vec![Some(gi)])
    })
}

impl Tensor {
    fn backend(&self) -> Result<Arc<dyn Backend>> {
        backend_for(self.device())
    }

    // ---- unary -----------------------------------------------------------

    fn unary_op(&self, op: UnaryOp) -> Result<Tensor> {
        let out = self.backend()?.unary(op, self)?;
        let a = self.clone();
        let o = out.clone();
        Ok(record(out, vec![self.clone()], move |g| {
            let gi = match op {
                UnaryOp::Neg => g.neg()?,
                UnaryOp::Exp => g.mul(&o)?,
                UnaryOp::Ln => g.div(&a)?,
                UnaryOp::Abs => {
                    let sign = a
                        .gt_mask(&Tensor::scalar_on(&a, 0.0)?)?
                        .sub(&Tensor::scalar_on(&a, 0.0)?.gt_mask(&a)?)?;
                    g.mul(&sign)?
                }
                UnaryOp::Sqrt => g.mul(&Tensor::scalar_on(&a, 0.5)?)?.div(&o)?,
                UnaryOp::Sin => g.mul(&a.cos()?)?,
                UnaryOp::Cos => g.mul(&a.sin()?.neg()?)?,
                UnaryOp::Tanh => {
                    let one = Tensor::scalar_on(&a, 1.0)?;
                    g.mul(&one.sub(&o.mul(&o)?)?)?
                }
                UnaryOp::Relu => g.mul(&a.gt_mask(&Tensor::scalar_on(&a, 0.0)?)?)?,
                UnaryOp::Gelu => {
                    // d/dx [0.5x(1+tanh(u))], u = c(x + 0.044715 x^3),
                    // c = sqrt(2/pi).
                    let c = Tensor::scalar_on(&a, 0.797_884_6)?;
                    let k = Tensor::scalar_on(&a, 0.044_715)?;
                    let one = Tensor::scalar_on(&a, 1.0)?;
                    let half = Tensor::scalar_on(&a, 0.5)?;
                    let three_k = Tensor::scalar_on(&a, 3.0 * 0.044_715)?;
                    let x2 = a.mul(&a)?;
                    let u = c.mul(&a.add(&k.mul(&x2.mul(&a)?)?)?)?;
                    let t = u.tanh()?;
                    let sech2 = one.sub(&t.mul(&t)?)?;
                    let du = c.mul(&one.add(&three_k.mul(&x2)?)?)?;
                    let d = half
                        .mul(&one.add(&t)?)?
                        .add(&half.mul(&a)?.mul(&sech2)?.mul(&du)?)?;
                    g.mul(&d)?
                }
                UnaryOp::Sigmoid => {
                    let one = Tensor::scalar_on(&a, 1.0)?;
                    g.mul(&o)?.mul(&one.sub(&o)?)?
                }
            };
            Ok(vec![Some(gi)])
        }))
    }

    /// Elementwise negation.
    pub fn neg(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Neg)
    }
    /// Elementwise `e^x`.
    pub fn exp(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Exp)
    }
    /// Elementwise natural logarithm.
    pub fn ln(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Ln)
    }
    /// Elementwise absolute value.
    pub fn abs(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Abs)
    }
    /// Elementwise square root.
    pub fn sqrt(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Sqrt)
    }
    /// Elementwise sine.
    pub fn sin(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Sin)
    }
    /// Elementwise cosine.
    pub fn cos(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Cos)
    }
    /// Elementwise hyperbolic tangent.
    pub fn tanh(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Tanh)
    }
    /// Elementwise rectified linear unit.
    pub fn relu(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Relu)
    }
    /// Elementwise GELU (tanh approximation).
    pub fn gelu(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Gelu)
    }
    /// Elementwise logistic sigmoid.
    pub fn sigmoid(&self) -> Result<Tensor> {
        self.unary_op(UnaryOp::Sigmoid)
    }

    /// A scalar constant with the dtype and device of `like` (plumbing for
    /// VJPs and scalar operator overloads).
    pub fn scalar_on(like: &Tensor, value: f32) -> Result<Tensor> {
        let s = match like.dtype() {
            DType::F64 => Tensor::from_vec_f64(vec![f64::from(value)], Shape::from([]))?,
            _ => Tensor::scalar(value),
        };
        s.to_device(like.device())
    }

    /// This tensor's elements converted to `dtype` (`F32` ↔ `F64`, or
    /// `I64` → float). A no-op clone for the same dtype. CPU only for
    /// `F64`; differentiable (the gradient converts back).
    pub fn to_dtype(&self, dtype: DType) -> Result<Tensor> {
        if self.dtype() == dtype {
            return Ok(self.clone());
        }
        if self.device() != Device::Cpu {
            return Err(Error::UnsupportedDType {
                dtype,
                op: "to_dtype (device tensors are f32; convert on the CPU)",
            });
        }
        let shape = self.shape().clone();
        let out = match (self.dtype(), dtype) {
            (DType::F32, DType::F64) => Tensor::from_vec_f64(
                self.to_vec_f32()?.into_iter().map(f64::from).collect(),
                shape,
            )?,
            (DType::F64, DType::F32) => Tensor::from_vec_f32(
                self.to_vec_f64()?.into_iter().map(|x| x as f32).collect(),
                shape,
            )?,
            (DType::I64, DType::F32) => Tensor::from_vec_f32(
                self.to_vec_i64()?.into_iter().map(|x| x as f32).collect(),
                shape,
            )?,
            (DType::I64, DType::F64) => Tensor::from_vec_f64(
                self.to_vec_i64()?.into_iter().map(|x| x as f64).collect(),
                shape,
            )?,
            (_, to) => {
                return Err(Error::UnsupportedDType {
                    dtype: to,
                    op: "to_dtype",
                });
            }
        };
        let from = self.dtype();
        Ok(record(out, vec![self.clone()], move |g| {
            Ok(vec![Some(g.to_dtype(from)?)])
        }))
    }

    // ---- binary ----------------------------------------------------------

    fn binary_op(&self, op: BinaryOp, rhs: &Tensor) -> Result<Tensor> {
        let device = same_device(self, rhs, "binary")?;
        let out = backend_for(device)?.binary(op, self, rhs)?;
        let (a, b) = (self.clone(), rhs.clone());
        let o = out.clone();
        Ok(record(out, vec![self.clone(), rhs.clone()], move |g| {
            let (ga, gb): (Option<Tensor>, Option<Tensor>) = match op {
                BinaryOp::Add => (Some(g.clone()), Some(g.clone())),
                BinaryOp::Sub => (Some(g.clone()), Some(g.neg()?)),
                BinaryOp::Mul => (Some(g.mul_raw(&b)?), Some(g.mul_raw(&a)?)),
                BinaryOp::Div => {
                    let ga = g.div_raw(&b)?;
                    let gb = g.mul_raw(&o)?.div_raw(&b)?.neg()?;
                    (Some(ga), Some(gb))
                }
                BinaryOp::Pow => {
                    let one = Tensor::scalar_on(&a, 1.0)?;
                    let ga = g.mul_raw(&b)?.mul_raw(&a.pow(&b.sub(&one)?)?)?;
                    let gb = g.mul_raw(&o)?.mul_raw(&a.ln()?)?;
                    (Some(ga), Some(gb))
                }
                // A tie splits the gradient evenly — the convention
                // docs/LIMITATIONS.md states. Sending all of it to one
                // operand made every composite built on `maximum` silently
                // wrong at a tie: BCEWithLogitsLoss is `max(x, 0) - x*t + …`,
                // so at a logit of exactly 0 it returned `-t` instead of
                // `sigmoid(0) - t`, and descent moved uphill.
                BinaryOp::Maximum => {
                    let half = a.eq_mask(&b)?.mul_scalar(0.5)?;
                    let wa = a.gt_mask(&b)?.add(&half)?;
                    let one = Tensor::scalar_on(&a, 1.0)?;
                    let ga = g.mul_raw(&wa)?;
                    let gb = g.mul_raw(&one.sub(&wa)?)?;
                    (Some(ga), Some(gb))
                }
                BinaryOp::Minimum => {
                    let half = a.eq_mask(&b)?.mul_scalar(0.5)?;
                    let wa = b.gt_mask(&a)?.add(&half)?;
                    let one = Tensor::scalar_on(&a, 1.0)?;
                    let ga = g.mul_raw(&wa)?;
                    let gb = g.mul_raw(&one.sub(&wa)?)?;
                    (Some(ga), Some(gb))
                }
                BinaryOp::Gt | BinaryOp::Eq => (None, None),
            };
            let ga = match ga {
                Some(t) => Some(reduce_to_shape(&t, a.shape())?),
                None => None,
            };
            let gb = match gb {
                Some(t) => Some(reduce_to_shape(&t, b.shape())?),
                None => None,
            };
            Ok(vec![ga, gb])
        }))
    }

    /// Untracked multiply, for use inside VJP closures (recording is
    /// already off during backward; this is belt and braces).
    fn mul_raw(&self, rhs: &Tensor) -> Result<Tensor> {
        let device = same_device(self, rhs, "mul")?;
        backend_for(device)?.binary(BinaryOp::Mul, self, rhs)
    }

    fn div_raw(&self, rhs: &Tensor) -> Result<Tensor> {
        let device = same_device(self, rhs, "div")?;
        backend_for(device)?.binary(BinaryOp::Div, self, rhs)
    }

    /// Elementwise addition, broadcasting.
    pub fn add(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Add, rhs)
    }
    /// Elementwise subtraction, broadcasting.
    pub fn sub(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Sub, rhs)
    }
    /// Elementwise multiplication, broadcasting.
    pub fn mul(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Mul, rhs)
    }
    /// Elementwise division, broadcasting.
    pub fn div(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Div, rhs)
    }
    /// Elementwise power, broadcasting.
    pub fn pow(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Pow, rhs)
    }
    /// Elementwise maximum, broadcasting.
    pub fn maximum(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Maximum, rhs)
    }
    /// Elementwise minimum, broadcasting.
    pub fn minimum(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Minimum, rhs)
    }
    /// Elementwise `a > b` as a 0.0/1.0 mask. Not differentiable.
    pub fn gt_mask(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Gt, rhs)
    }
    /// Elementwise `a == b` as a 0.0/1.0 mask. Not differentiable.
    pub fn eq_mask(&self, rhs: &Tensor) -> Result<Tensor> {
        self.binary_op(BinaryOp::Eq, rhs)
    }

    /// Add a scalar, broadcasting.
    pub fn add_scalar(&self, s: f32) -> Result<Tensor> {
        self.add(&Tensor::scalar_on(self, s)?)
    }
    /// Multiply by a scalar, broadcasting.
    pub fn mul_scalar(&self, s: f32) -> Result<Tensor> {
        self.mul(&Tensor::scalar_on(self, s)?)
    }

    // ---- matmul ----------------------------------------------------------

    /// Matrix product with NumPy/PyTorch batch semantics.
    ///
    /// The last two dimensions are the matrix (`[.., m, k] x [.., k, n]`
    /// → `[.., m, n]`); every leading dimension is a batch dimension, and
    /// batch dimensions broadcast against each other (1 against `b`, and
    /// a missing leading dimension counts as 1). A rank-2 operand is one
    /// matrix for every batch of the other. The result is rank 2 only
    /// when both operands are: `[2, 2, 3] x [1, 3, 2]` is `[2, 2, 2]`,
    /// `[m, k] x [b, k, n]` is `[b, m, n]`, `[2, 1, 3, 4] x [5, 4, 6]` is
    /// `[2, 5, 3, 6]`.
    ///
    /// Backends implement the rank-2/rank-3 contract of
    /// [`plan_matmul`](crate::backend::plan_matmul); higher ranks are
    /// lowered here — the batch dimensions are broadcast (a zero-stride
    /// view, materialized only when an operand's batch really has to be
    /// repeated), flattened to one batch axis, multiplied, and unflattened
    /// — and every step is a recorded op, so the gradient needs no VJP of
    /// its own.
    pub fn matmul(&self, rhs: &Tensor) -> Result<Tensor> {
        let device = same_device(self, rhs, "matmul")?;
        if self.ndim() > 3 || rhs.ndim() > 3 {
            return self.matmul_lowered(rhs);
        }
        let out = backend_for(device)?.matmul(self, rhs)?;
        let (a, b) = (self.clone(), rhs.clone());
        Ok(record(out, vec![self.clone(), rhs.clone()], move |g| {
            // g is [.., m, n]; the operand grads carry g's batch, which is
            // then summed down onto a broadcast (or rank-2) operand.
            let ga = backend_for(g.device())?.matmul(g, &b.t()?)?;
            let gb = backend_for(g.device())?.matmul(&a.t()?, g)?;
            Ok(vec![
                Some(reduce_to_shape(&ga, a.shape())?),
                Some(reduce_to_shape(&gb, b.shape())?),
            ])
        }))
    }

    /// Rank ≥ 4 matmul: broadcast the batch dimensions, flatten them to one,
    /// run the rank-3 contract, unflatten. Composed from recorded view ops.
    fn matmul_lowered(&self, rhs: &Tensor) -> Result<Tensor> {
        let (ad, bd) = (self.dims(), rhs.dims());
        if ad.len() < 2 || bd.len() < 2 {
            return Err(Error::InvalidArgument {
                op: "matmul",
                detail: format!("operands need rank >= 2; got {}x{}", ad.len(), bd.len()),
            });
        }
        let (m, k) = (ad[ad.len() - 2], ad[ad.len() - 1]);
        let (kb, n) = (bd[bd.len() - 2], bd[bd.len() - 1]);
        if k != kb {
            return Err(Error::ShapeMismatch {
                expected: Shape::new(bd[..bd.len() - 2].iter().copied().chain([k, n]).collect()),
                got: rhs.shape().clone(),
                op: "matmul",
            });
        }
        let a_batch = Shape::new(ad[..ad.len() - 2].to_vec());
        let b_batch = Shape::new(bd[..bd.len() - 2].to_vec());
        let batch = oxmera_core::shape::broadcast_shapes(&a_batch, &b_batch).map_err(|_| {
            Error::BroadcastIncompatible {
                lhs: self.shape().clone(),
                rhs: rhs.shape().clone(),
            }
        })?;
        let batch_numel = batch.numel();
        // An operand whose batch is a single matrix stays rank 2 and lets
        // the backend broadcast it with a zero batch stride; anything else
        // is expanded to the full batch and flattened.
        let lower = |t: &Tensor, own: &Shape, rows: usize, cols: usize| -> Result<Tensor> {
            if own.numel() == 1 {
                return t.reshape(Shape::from([rows, cols]));
            }
            let full: Vec<usize> = batch.dims().iter().copied().chain([rows, cols]).collect();
            let expanded = if own.dims() == batch.dims() {
                t.clone()
            } else {
                // Right-align the operand's batch dims under the broadcast
                // batch, then take the (recorded) zero-stride view.
                let lead = batch.ndim() - own.ndim();
                let padded: Vec<usize> = std::iter::repeat_n(1usize, lead)
                    .chain(own.dims().iter().copied())
                    .chain([rows, cols])
                    .collect();
                t.reshape(Shape::new(padded))?
                    .broadcast_to(Shape::new(full.clone()))?
                    .contiguous()?
            };
            expanded.reshape(Shape::from([batch_numel, rows, cols]))
        };
        let a3 = lower(self, &a_batch, m, k)?;
        let b3 = lower(rhs, &b_batch, k, n)?;
        let out = a3.matmul(&b3)?;
        let out_shape: Vec<usize> = batch.dims().iter().copied().chain([m, n]).collect();
        out.reshape(Shape::new(out_shape))
    }

    // ---- reductions --------------------------------------------------------

    fn reduce_op(&self, op: ReduceOp, axes: &[usize], keepdim: bool) -> Result<Tensor> {
        let axes = normalize_axes(axes, self.ndim(), "reduce")?;
        let out = self.backend()?.reduce(op, self, &axes, keepdim)?;
        let a = self.clone();
        let o = out.clone();
        let axes_c = axes.clone();
        Ok(record(out, vec![self.clone()], move |g| {
            // Re-insert reduced axes as size 1 so broadcasting lines up.
            let g_keep = if keepdim {
                g.clone()
            } else {
                unsqueeze_axes(g, &axes_c)?
            };
            let gi = match op {
                ReduceOp::Sum => g_keep.broadcast_to(a.shape().clone())?.contiguous()?,
                ReduceOp::Max | ReduceOp::Min => {
                    let o_keep = if keepdim {
                        o.clone()
                    } else {
                        unsqueeze_axes(&o, &axes_c)?
                    };
                    let mask = a.eq_mask(&o_keep.broadcast_to(a.shape().clone())?)?;
                    let count = mask.sum_keepdim(&axes_c, true)?;
                    g_keep
                        .broadcast_to(a.shape().clone())?
                        .mul_raw(&mask)?
                        .div_raw(&count.broadcast_to(a.shape().clone())?.contiguous()?)?
                }
            };
            Ok(vec![Some(gi)])
        }))
    }

    /// Sum over `axes` (empty means all), removing them from the shape.
    pub fn sum(&self, axes: &[usize]) -> Result<Tensor> {
        self.reduce_op(ReduceOp::Sum, axes, false)
    }

    /// Sum over `axes` with explicit `keepdim`.
    pub fn sum_keepdim(&self, axes: &[usize], keepdim: bool) -> Result<Tensor> {
        self.reduce_op(ReduceOp::Sum, axes, keepdim)
    }

    /// Maximum over `axes` (empty means all).
    pub fn max(&self, axes: &[usize]) -> Result<Tensor> {
        self.reduce_op(ReduceOp::Max, axes, false)
    }

    /// Maximum over `axes` with explicit `keepdim`.
    pub fn max_keepdim(&self, axes: &[usize], keepdim: bool) -> Result<Tensor> {
        self.reduce_op(ReduceOp::Max, axes, keepdim)
    }

    /// Minimum over `axes` (empty means all).
    pub fn min(&self, axes: &[usize]) -> Result<Tensor> {
        self.reduce_op(ReduceOp::Min, axes, false)
    }

    /// Mean over `axes` (empty means all) — composite, so its gradient
    /// flows through `sum` and scalar multiply.
    ///
    /// Reducing over a zero-length extent is a typed error; see
    /// [`Tensor::mean_keepdim`].
    pub fn mean(&self, axes: &[usize]) -> Result<Tensor> {
        self.mean_keepdim(axes, false)
    }

    /// Mean over `axes` with explicit `keepdim`.
    ///
    /// # Errors
    ///
    /// Reducing over a zero-length extent is an
    /// [`Error::InvalidArgument`]: the mean of nothing is `0/0`, and there
    /// is no value that is the right answer.
    ///
    /// Until 0.4.0 this returned `NaN`, which is not an answer either but
    /// looks like one. A `NaN` from an empty last batch does not fail — it
    /// flows into the loss, then into every gradient, and surfaces an
    /// epoch later as a model that stopped learning for no visible reason.
    /// `argmax` already refused the same input for the same reason; this
    /// is the pair being made consistent.
    ///
    /// `sum` and `max` still return their identities (`0` and `-inf`) over
    /// an empty extent, and deliberately: those compose correctly under
    /// further reduction, and `mean` does not. The whole family is
    /// tabulated in `docs/LIMITATIONS.md`.
    pub fn mean_keepdim(&self, axes: &[usize], keepdim: bool) -> Result<Tensor> {
        let axes_n = normalize_axes(axes, self.ndim(), "mean")?;
        let n: usize = axes_n.iter().map(|&ax| self.dims()[ax]).product();
        if n == 0 {
            let empty: Vec<usize> = axes_n
                .iter()
                .copied()
                .filter(|&ax| self.dims()[ax] == 0)
                .collect();
            return Err(Error::InvalidArgument {
                op: "mean",
                detail: format!(
                    "dimension(s) {empty:?} have extent 0, so the mean would be 0/0; \
                     the mean of nothing is undefined — use sum() if an empty \
                     reduction should be 0, or guard the empty case at the call site"
                ),
            });
        }
        let summed = self.sum_keepdim(&axes_n, keepdim)?;
        if summed.dtype() == DType::F64 {
            // Scaling an f64 sum by an f32 reciprocal threw away every digit
            // past f32 precision: mean([1,2,3]) came back 2.0000000596046448.
            // Divide in the tensor's own dtype instead.
            let divisor = Tensor::from_vec_f64(vec![n as f64], Shape::from([]))?
                .to_device(summed.device())?;
            return summed.div(&divisor);
        }
        summed.mul_scalar(1.0 / n as f32)
    }

    /// Index of the maximum along `dim`, as an `I64` tensor. Not
    /// differentiable.
    pub fn argmax(&self, dim: usize, keepdim: bool) -> Result<Tensor> {
        if dim >= self.ndim() {
            return Err(Error::InvalidArgument {
                op: "argmax",
                detail: format!("dim {dim} out of range for rank {}", self.ndim()),
            });
        }
        self.backend()?.argmax(self, dim, keepdim)
    }

    /// Numerically stable softmax along `dim` — composite.
    pub fn softmax(&self, dim: usize) -> Result<Tensor> {
        // The row max and row sum stay broadcast views: the binary kernels
        // consume a stride-0 operand directly, so nothing is materialized.
        let shifted = self.sub(&self.max_keepdim(&[dim], true)?.detach())?;
        let e = shifted.exp()?;
        let denom = e.sum_keepdim(&[dim], true)?;
        e.div(&denom)
    }

    /// Numerically stable log-softmax along `dim` — composite.
    pub fn log_softmax(&self, dim: usize) -> Result<Tensor> {
        let shifted = self.sub(&self.max_keepdim(&[dim], true)?.detach())?;
        let lse = shifted.exp()?.sum_keepdim(&[dim], true)?.ln()?;
        shifted.sub(&lse)
    }

    // ---- indexing -----------------------------------------------------------

    /// Rows of `self` along `dim` selected by `indices` (`I64`, on the CPU).
    pub fn index_select(&self, dim: usize, indices: &Tensor) -> Result<Tensor> {
        let out = dispatch_index(self, &[indices], |be, t, extra| {
            be.index_select(t, dim, &extra[0])
        })?;
        let in_shape = self.shape().clone();
        let idx = indices.clone();
        Ok(record(out, vec![self.clone()], move |g| {
            let zeros = Tensor::zeros(in_shape.clone())
                .to_dtype(g.dtype())?
                .to_device(g.device())?;
            Ok(vec![Some(zeros.index_add(dim, &idx, g)?)])
        }))
    }

    /// `out[indices[i]] += src[i]` along `dim`, on a fresh copy of `self`.
    /// `indices` is `I64` on the CPU; `src` lives on `self`'s device.
    pub fn index_add(&self, dim: usize, indices: &Tensor, src: &Tensor) -> Result<Tensor> {
        same_device(self, src, "index_add")?;
        let out = dispatch_index(self, &[indices, src], |be, t, extra| {
            be.index_add(t, dim, &extra[0], &extra[1])
        })?;
        let idx = indices.clone();
        Ok(record(out, vec![self.clone(), src.clone()], move |g| {
            Ok(vec![Some(g.clone()), Some(g.index_select(dim, &idx)?)])
        }))
    }

    // ---- device movement ------------------------------------------------------

    /// This tensor's data on `device` (a cheap clone when already there).
    /// `F64` tensors are CPU-only: moving one to a GPU is a typed error.
    pub fn to_device(&self, device: Device) -> Result<Tensor> {
        if self.device() == device {
            return Ok(self.clone());
        }
        if self.dtype() == DType::F64 {
            return Err(Error::UnsupportedDType {
                dtype: DType::F64,
                op: "to_device (f64 tensors live on the CPU; to_dtype(F32) first)",
            });
        }
        let out = match (self.device(), device) {
            (Device::Cpu, target) => backend_for(target)?.upload(&self.contiguous_data()?)?,
            (_, Device::Cpu) => self.backend()?.download(self)?,
            (_, target) => {
                let host = self.backend()?.download(self)?;
                backend_for(target)?.upload(&host)?
            }
        };
        let source = self.device();
        Ok(record(out, vec![self.clone()], move |g| {
            Ok(vec![Some(g.to_device(source)?)])
        }))
    }
}

/// Insert size-1 axes at `axes` (sorted ascending) — plumbing for reduce
/// VJPs.
fn unsqueeze_axes(t: &Tensor, axes: &[usize]) -> Result<Tensor> {
    let mut out = t.clone();
    let mut sorted = axes.to_vec();
    sorted.sort_unstable();
    for &ax in &sorted {
        out = out.unsqueeze(ax)?;
    }
    Ok(out)
}

/// Validate and canonicalize reduce axes; empty means all axes.
fn normalize_axes(axes: &[usize], ndim: usize, op: &'static str) -> Result<Vec<usize>> {
    let mut axes: Vec<usize> = if axes.is_empty() {
        (0..ndim).collect()
    } else {
        axes.to_vec()
    };
    axes.sort_unstable();
    axes.dedup();
    if let Some(&bad) = axes.iter().find(|&&a| a >= ndim) {
        return Err(Error::InvalidArgument {
            op,
            detail: format!("axis {bad} out of range for rank {ndim}"),
        });
    }
    Ok(axes)
}

/// Run an index op on the tensor's backend, falling back to a CPU
/// round-trip when the backend declines. Every tensor operand — the
/// indices and, for `index_add`, the source — takes the round-trip too:
/// moving only `t` left `src` on the device and failed the CPU backend
/// with a DeviceMismatch inside the `narrow` VJP (found by oxmega's
/// k-DPP loss on CUDA).
fn dispatch_index(
    t: &Tensor,
    extra: &[&Tensor],
    f: impl Fn(&dyn Backend, &Tensor, &[Tensor]) -> Result<Tensor>,
) -> Result<Tensor> {
    let backend = backend_for(t.device())?;
    let on_device: Vec<Tensor> = extra.iter().map(|e| (*e).clone()).collect();
    match f(backend.as_ref(), t, &on_device) {
        Err(Error::NotImplemented { .. }) if t.device() != Device::Cpu => {
            let cpu = backend.download(t)?;
            let cpu_extra: Vec<Tensor> = extra
                .iter()
                .map(|e| e.to_device(Device::Cpu))
                .collect::<Result<_>>()?;
            let cpu_backend = backend_for(Device::Cpu)?;
            let out = f(cpu_backend.as_ref(), &cpu, &cpu_extra)?;
            backend_for(t.device())?.upload(&out)
        }
        other => other,
    }
}