mlx_core/

tensor.rs

//! Tensor type — a lazy handle to a node in the computation graph.
//!
//! Operations on tensors record nodes in the graph. Actual computation is
//! deferred until `eval()` (or `to_vec_f32()`) is called, at which point the
//! stream topologically sorts the subgraph and dispatches to the backend.

use std::sync::Arc;

use smallvec::SmallVec;

use crate::backend::{Stream, default_stream};
use crate::graph::{OpKind, TensorMeta};
use crate::{DType, MlxError, NodeId, Result, Shape};

/// Compute device.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Device {
    Cpu,
    Gpu,
}

impl Device {
    /// Return the default device for this platform.
    pub fn default_device() -> Self {
        #[cfg(target_os = "macos")]
        {
            // TODO: Enable Metal backend when available.
            Device::Cpu
        }
        #[cfg(not(target_os = "macos"))]
        {
            Device::Cpu
        }
    }
}

/// A tensor handle.
///
/// In the lazy graph model a `Tensor` is a lightweight reference to a node in
/// the computation graph. Operations build up the graph; actual computation
/// happens when `eval()` is called (or implicitly via `to_vec_f32()`).
#[derive(Clone)]
pub struct Tensor {
    node_id: NodeId,
    shape: Shape,
    dtype: DType,
    device: Device,
    stream: Arc<Stream>,
}

impl Tensor {
    // ── Constructors ────────────────────────────────────────────────────

    /// Create a tensor filled with zeros.
    pub fn zeros(shape: &Shape, dtype: DType, device: &Device) -> Result<Self> {
        let n = shape.numel() as usize;
        Self::from_data(vec![0.0; n], shape, dtype, device)
    }

    /// Create a tensor filled with ones.
    pub fn ones(shape: &Shape, dtype: DType, device: &Device) -> Result<Self> {
        let n = shape.numel() as usize;
        Self::from_data(vec![1.0; n], shape, dtype, device)
    }

    /// Create a tensor from f32 data.
    pub fn from_f32(data: &[f32], shape: &Shape, device: &Device) -> Result<Self> {
        let expected = shape.numel() as usize;
        if data.len() != expected {
            return Err(MlxError::InvalidArgument(format!(
                "data length {} does not match shape {} (expected {})",
                data.len(),
                shape,
                expected,
            )));
        }
        Self::from_data(data.to_vec(), shape, DType::F32, device)
    }

    /// Create a tensor from f32 data on a specific stream.
    pub fn from_f32_on_stream(data: &[f32], shape: &Shape, stream: &Arc<Stream>) -> Result<Self> {
        let expected = shape.numel() as usize;
        if data.len() != expected {
            return Err(MlxError::InvalidArgument(format!(
                "data length {} does not match shape {} (expected {})",
                data.len(),
                shape,
                expected,
            )));
        }
        let meta = TensorMeta {
            shape: shape.clone(),
            dtype: DType::F32,
        };
        let node_id = stream.add_constant(data.to_vec(), meta);
        Ok(Self {
            node_id,
            shape: shape.clone(),
            dtype: DType::F32,
            device: Device::Gpu, // Assume GPU if custom stream for now, or detect
            stream: Arc::clone(stream),
        })
    }

    /// Create a tensor from f32 data with a specified dtype.
    ///
    /// The data is stored as f32 internally; `dtype` records the logical type
    /// (e.g. F16 weights that were converted to f32 on load).
    pub fn from_data_with_dtype(
        data: Vec<f32>,
        shape: &Shape,
        dtype: DType,
        device: &Device,
    ) -> Result<Self> {
        let expected = shape.numel() as usize;
        if data.len() != expected {
            return Err(MlxError::InvalidArgument(format!(
                "data length {} does not match shape {} (expected {})",
                data.len(),
                shape,
                expected,
            )));
        }
        Self::from_data(data, shape, dtype, device)
    }

    fn from_data(data: Vec<f32>, shape: &Shape, dtype: DType, device: &Device) -> Result<Self> {
        let stream = default_stream();
        let meta = TensorMeta {
            shape: shape.clone(),
            dtype,
        };
        let node_id = stream.add_constant(data, meta);
        Ok(Self {
            node_id,
            shape: shape.clone(),
            dtype,
            device: device.clone(),
            stream,
        })
    }

    fn lazy_op(
        &self,
        op: OpKind,
        inputs: SmallVec<[NodeId; 2]>,
        shape: Shape,
        dtype: DType,
    ) -> Self {
        let meta = TensorMeta {
            shape: shape.clone(),
            dtype,
        };
        let node_id = self.stream.add_op(op, inputs, meta);
        Tensor {
            node_id,
            shape,
            dtype,
            device: self.device.clone(),
            stream: Arc::clone(&self.stream),
        }
    }

    // ── Elementwise ops ─────────────────────────────────────────────────

    /// Element-wise addition.
    pub fn add(&self, rhs: &Tensor) -> Result<Tensor> {
        if self.shape != rhs.shape {
            return Err(MlxError::ShapeMismatch {
                expected: self.shape.0.clone(),
                got: rhs.shape.0.clone(),
            });
        }
        Ok(self.lazy_op(
            OpKind::Add,
            SmallVec::from_slice(&[self.node_id, rhs.node_id]),
            self.shape.clone(),
            self.dtype,
        ))
    }

    /// Element-wise subtraction.
    pub fn sub(&self, rhs: &Tensor) -> Result<Tensor> {
        if self.shape != rhs.shape {
            return Err(MlxError::ShapeMismatch {
                expected: self.shape.0.clone(),
                got: rhs.shape.0.clone(),
            });
        }
        Ok(self.lazy_op(
            OpKind::Sub,
            SmallVec::from_slice(&[self.node_id, rhs.node_id]),
            self.shape.clone(),
            self.dtype,
        ))
    }

    /// Element-wise multiplication.
    pub fn mul(&self, rhs: &Tensor) -> Result<Tensor> {
        if self.shape != rhs.shape {
            return Err(MlxError::ShapeMismatch {
                expected: self.shape.0.clone(),
                got: rhs.shape.0.clone(),
            });
        }
        Ok(self.lazy_op(
            OpKind::Mul,
            SmallVec::from_slice(&[self.node_id, rhs.node_id]),
            self.shape.clone(),
            self.dtype,
        ))
    }

    /// Element-wise division.
    pub fn div(&self, rhs: &Tensor) -> Result<Tensor> {
        if self.shape != rhs.shape {
            return Err(MlxError::ShapeMismatch {
                expected: self.shape.0.clone(),
                got: rhs.shape.0.clone(),
            });
        }
        Ok(self.lazy_op(
            OpKind::Div,
            SmallVec::from_slice(&[self.node_id, rhs.node_id]),
            self.shape.clone(),
            self.dtype,
        ))
    }

    /// Element-wise negation.
    pub fn neg(&self) -> Tensor {
        self.lazy_op(
            OpKind::Neg,
            SmallVec::from_slice(&[self.node_id]),
            self.shape.clone(),
            self.dtype,
        )
    }

    // ── Reductions ──────────────────────────────────────────────────────

    /// Sum along an axis.
    pub fn sum_axis(&self, axis: i32) -> Result<Tensor> {
        let ndim = self.shape.ndim() as i32;
        let ax = if axis < 0 { ndim + axis } else { axis };
        if ax < 0 || ax >= ndim {
            return Err(MlxError::InvalidArgument(format!(
                "axis {axis} out of range for ndim {ndim}"
            )));
        }
        let mut new_dims: Vec<i64> = self.shape.0.clone();
        new_dims.remove(ax as usize);
        Ok(self.lazy_op(
            OpKind::Sum { axis: Some(ax) },
            SmallVec::from_slice(&[self.node_id]),
            Shape::new(new_dims),
            self.dtype,
        ))
    }

    /// Sum all elements to a scalar.
    pub fn sum_all(&self) -> Result<Tensor> {
        Ok(self.lazy_op(
            OpKind::Sum { axis: None },
            SmallVec::from_slice(&[self.node_id]),
            Shape::scalar(),
            self.dtype,
        ))
    }

    // ── Linear algebra ──────────────────────────────────────────────────

    /// Matrix multiplication (2D only for now).
    pub fn matmul(&self, rhs: &Tensor) -> Result<Tensor> {
        if self.shape.ndim() != 2 || rhs.shape.ndim() != 2 {
            return Err(MlxError::InvalidArgument(
                "matmul requires 2D tensors".to_string(),
            ));
        }
        let m = self.shape.0[0];
        let k = self.shape.0[1];
        let k2 = rhs.shape.0[0];
        let n = rhs.shape.0[1];
        if k != k2 {
            return Err(MlxError::ShapeMismatch {
                expected: self.shape.0.clone(),
                got: rhs.shape.0.clone(),
            });
        }
        Ok(self.lazy_op(
            OpKind::MatMul,
            SmallVec::from_slice(&[self.node_id, rhs.node_id]),
            Shape::new(vec![m, n]),
            self.dtype,
        ))
    }

    // ── Shape manipulation ──────────────────────────────────────────────

    /// Reshape the tensor.
    pub fn reshape(&self, new_shape: &Shape) -> Result<Tensor> {
        if self.shape.numel() != new_shape.numel() {
            return Err(MlxError::ShapeMismatch {
                expected: self.shape.0.clone(),
                got: new_shape.0.clone(),
            });
        }
        Ok(self.lazy_op(
            OpKind::Reshape {
                new_shape: new_shape.clone(),
            },
            SmallVec::from_slice(&[self.node_id]),
            new_shape.clone(),
            self.dtype,
        ))
    }

    /// Transpose (reverses axes by default, or use specified permutation).
    pub fn transpose(&self, axes: Option<&[usize]>) -> Result<Tensor> {
        let ndim = self.shape.ndim();
        let perm: Vec<usize> = match axes {
            Some(ax) => {
                if ax.len() != ndim {
                    return Err(MlxError::InvalidArgument(
                        "transpose axes length must match ndim".into(),
                    ));
                }
                let mut seen = vec![false; ndim];
                for &axis in ax {
                    if axis >= ndim {
                        return Err(MlxError::InvalidArgument(format!(
                            "transpose axis {axis} out of range for ndim {ndim}"
                        )));
                    }
                    if seen[axis] {
                        return Err(MlxError::InvalidArgument(format!(
                            "duplicate transpose axis {axis} in axes; expected a permutation of 0..{ndim}"
                        )));
                    }
                    seen[axis] = true;
                }
                ax.to_vec()
            }
            None => (0..ndim).rev().collect(),
        };
        let new_dims: Vec<i64> = perm.iter().map(|&ax| self.shape.0[ax]).collect();
        Ok(self.lazy_op(
            OpKind::Transpose { axes: Some(perm) },
            SmallVec::from_slice(&[self.node_id]),
            Shape::new(new_dims),
            self.dtype,
        ))
    }

    // ── Activations ─────────────────────────────────────────────────────

    /// Softmax along an axis.
    pub fn softmax(&self, axis: i32) -> Result<Tensor> {
        let ndim = self.shape.ndim() as i32;
        let ax = if axis < 0 { ndim + axis } else { axis };
        if ax < 0 || ax >= ndim {
            return Err(MlxError::InvalidArgument(format!(
                "axis {axis} out of range for ndim {ndim}"
            )));
        }
        Ok(self.lazy_op(
            OpKind::Softmax { axis },
            SmallVec::from_slice(&[self.node_id]),
            self.shape.clone(),
            self.dtype,
        ))
    }

    /// SiLU (Sigmoid Linear Unit) activation.
    pub fn silu(&self) -> Tensor {
        self.lazy_op(
            OpKind::Silu,
            SmallVec::from_slice(&[self.node_id]),
            self.shape.clone(),
            self.dtype,
        )
    }

    /// GELU (Gaussian Error Linear Unit) activation.
    pub fn gelu(&self) -> Tensor {
        self.lazy_op(
            OpKind::Gelu,
            SmallVec::from_slice(&[self.node_id]),
            self.shape.clone(),
            self.dtype,
        )
    }

    // ── Normalization ───────────────────────────────────────────────────

    /// Layer normalization over the last dimension.
    pub fn layer_norm(&self, eps: f32) -> Tensor {
        self.lazy_op(
            OpKind::LayerNorm { eps },
            SmallVec::from_slice(&[self.node_id]),
            self.shape.clone(),
            self.dtype,
        )
    }

    /// RMS normalization over the last dimension.
    pub fn rms_norm(&self, eps: f32) -> Tensor {
        self.lazy_op(
            OpKind::RmsNorm { eps },
            SmallVec::from_slice(&[self.node_id]),
            self.shape.clone(),
            self.dtype,
        )
    }

    /// Apply Rotary Positional Embeddings.
    pub fn rope(&self, rotary_dim: usize, pos_offset: usize, theta: f32) -> Tensor {
        self.lazy_op(
            OpKind::Rope {
                rotary_dim,
                pos_offset,
                theta,
            },
            SmallVec::from_slice(&[self.node_id]),
            self.shape.clone(),
            self.dtype,
        )
    }

    // ── Backward (VJP) helpers ─────────────────────────────────────────

    /// LayerNorm VJP: compute grad_input given grad_output and original input.
    pub fn layer_norm_vjp(&self, input: &Tensor, eps: f32) -> Result<Tensor> {
        if self.shape != input.shape {
            return Err(MlxError::ShapeMismatch {
                expected: input.shape.0.clone(),
                got: self.shape.0.clone(),
            });
        }
        if self.dtype != input.dtype {
            return Err(MlxError::InvalidArgument(
                "layer_norm_vjp requires matching dtypes".into(),
            ));
        }
        if self.device != input.device {
            return Err(MlxError::InvalidArgument(
                "layer_norm_vjp requires matching devices".into(),
            ));
        }
        Ok(self.lazy_op(
            OpKind::LayerNormVjp { eps },
            SmallVec::from_slice(&[self.node_id, input.node_id]),
            input.shape.clone(),
            input.dtype,
        ))
    }

    /// RmsNorm VJP: compute grad_input given grad_output and original input.
    pub fn rms_norm_vjp(&self, input: &Tensor, eps: f32) -> Result<Tensor> {
        if self.shape != input.shape {
            return Err(MlxError::ShapeMismatch {
                expected: input.shape.0.clone(),
                got: self.shape.0.clone(),
            });
        }
        if self.dtype != input.dtype {
            return Err(MlxError::InvalidArgument(
                "rms_norm_vjp requires matching dtypes".into(),
            ));
        }
        if self.device != input.device {
            return Err(MlxError::InvalidArgument(
                "rms_norm_vjp requires matching devices".into(),
            ));
        }
        Ok(self.lazy_op(
            OpKind::RmsNormVjp { eps },
            SmallVec::from_slice(&[self.node_id, input.node_id]),
            input.shape.clone(),
            input.dtype,
        ))
    }

    /// Softmax VJP: compute grad_input given grad_output (self) and softmax output.
    pub fn softmax_vjp(&self, softmax_output: &Tensor, axis: i32) -> Result<Tensor> {
        if self.shape != softmax_output.shape {
            return Err(MlxError::ShapeMismatch {
                expected: softmax_output.shape.0.clone(),
                got: self.shape.0.clone(),
            });
        }
        if self.dtype != softmax_output.dtype {
            return Err(MlxError::InvalidArgument(
                "softmax_vjp requires matching dtypes".into(),
            ));
        }
        if self.device != softmax_output.device {
            return Err(MlxError::InvalidArgument(
                "softmax_vjp requires matching devices".into(),
            ));
        }
        Ok(self.lazy_op(
            OpKind::SoftmaxVjp { axis },
            SmallVec::from_slice(&[self.node_id, softmax_output.node_id]),
            softmax_output.shape.clone(),
            softmax_output.dtype,
        ))
    }

    /// SiLU VJP: compute grad_input given grad_output (self) and original input.
    pub fn silu_vjp(&self, input: &Tensor) -> Result<Tensor> {
        if self.shape != input.shape {
            return Err(MlxError::ShapeMismatch {
                expected: input.shape.0.clone(),
                got: self.shape.0.clone(),
            });
        }
        if self.dtype != input.dtype {
            return Err(MlxError::InvalidArgument(
                "silu_vjp requires matching dtypes".into(),
            ));
        }
        if self.device != input.device {
            return Err(MlxError::InvalidArgument(
                "silu_vjp requires matching devices".into(),
            ));
        }
        Ok(self.lazy_op(
            OpKind::SiluVjp,
            SmallVec::from_slice(&[self.node_id, input.node_id]),
            input.shape.clone(),
            input.dtype,
        ))
    }

    /// GELU VJP: compute grad_input given grad_output (self) and original input.
    pub fn gelu_vjp(&self, input: &Tensor) -> Result<Tensor> {
        if self.shape != input.shape {
            return Err(MlxError::ShapeMismatch {
                expected: input.shape.0.clone(),
                got: self.shape.0.clone(),
            });
        }
        if self.dtype != input.dtype {
            return Err(MlxError::InvalidArgument(
                "gelu_vjp requires matching dtypes".into(),
            ));
        }
        if self.device != input.device {
            return Err(MlxError::InvalidArgument(
                "gelu_vjp requires matching devices".into(),
            ));
        }
        Ok(self.lazy_op(
            OpKind::GeluVjp,
            SmallVec::from_slice(&[self.node_id, input.node_id]),
            input.shape.clone(),
            input.dtype,
        ))
    }

    // ── Indexing / gathering ──────────────────────────────────────────

    /// Embedding lookup: gather rows from this weight matrix [vocab, dim]
    /// using `indices` [seq_len]. Returns [seq_len, dim].
    pub fn embedding_lookup(&self, indices: &Tensor) -> Result<Tensor> {
        if self.shape.ndim() != 2 {
            return Err(MlxError::InvalidArgument(
                "embedding_lookup: weight must be 2D [vocab_size, embed_dim]".into(),
            ));
        }
        if indices.shape.ndim() != 1 {
            return Err(MlxError::InvalidArgument(
                "embedding_lookup: indices must be 1D [seq_len]".into(),
            ));
        }
        let seq_len = indices.shape.0[0];
        let embed_dim = self.shape.0[1];
        Ok(self.lazy_op(
            OpKind::Embedding,
            SmallVec::from_slice(&[self.node_id, indices.node_id]),
            Shape::new(vec![seq_len, embed_dim]),
            self.dtype,
        ))
    }

    /// Narrow (slice) along an axis: extract `length` elements starting at `start`.
    pub fn narrow(&self, axis: i32, start: i64, length: i64) -> Result<Tensor> {
        let ndim = self.shape.ndim() as i32;
        let ax = if axis < 0 { ndim + axis } else { axis };
        if ax < 0 || ax >= ndim {
            return Err(MlxError::InvalidArgument(format!(
                "narrow: axis {axis} out of range for ndim {ndim}"
            )));
        }
        let ax_usize = ax as usize;
        let dim_size = self.shape.0[ax_usize];
        if start < 0 || start + length > dim_size {
            return Err(MlxError::InvalidArgument(format!(
                "narrow: start {start} + length {length} exceeds dim size {dim_size}"
            )));
        }
        let mut new_dims = self.shape.0.clone();
        new_dims[ax_usize] = length;
        Ok(self.lazy_op(
            OpKind::Narrow {
                axis: ax,
                start,
                length,
            },
            SmallVec::from_slice(&[self.node_id]),
            Shape::new(new_dims),
            self.dtype,
        ))
    }

    /// Concatenate tensors along an axis.
    pub fn cat(tensors: &[&Tensor], axis: i32) -> Result<Tensor> {
        if tensors.is_empty() {
            return Err(MlxError::InvalidArgument(
                "cat requires at least one tensor".into(),
            ));
        }
        let first = tensors[0];
        let ndim = first.shape.ndim() as i32;
        let ax = if axis < 0 { ndim + axis } else { axis };
        if ax < 0 || ax >= ndim {
            return Err(MlxError::InvalidArgument(format!(
                "cat: axis {axis} out of range for ndim {ndim}"
            )));
        }
        let ax_usize = ax as usize;

        // Validate shapes and compute output dim
        let mut total_dim: i64 = 0;
        for t in tensors {
            if t.shape.ndim() != first.shape.ndim() {
                return Err(MlxError::InvalidArgument(
                    "cat: all tensors must have same ndim".into(),
                ));
            }
            for (d, (&a, &b)) in first.shape.0.iter().zip(t.shape.0.iter()).enumerate() {
                if d != ax_usize && a != b {
                    return Err(MlxError::ShapeMismatch {
                        expected: first.shape.0.clone(),
                        got: t.shape.0.clone(),
                    });
                }
            }
            total_dim += t.shape.0[ax_usize];
        }

        let mut new_dims = first.shape.0.clone();
        new_dims[ax_usize] = total_dim;

        let inputs: SmallVec<[NodeId; 2]> = tensors.iter().map(|t| t.node_id).collect();

        Ok(first.lazy_op(
            OpKind::Concatenate { axis: ax },
            inputs,
            Shape::new(new_dims),
            first.dtype,
        ))
    }

    /// Single-head attention: Q @ K^T * scale → causal mask → softmax → @ V.
    /// Q: [Tq, Dh], K: [Tk, Dh], V: [Tk, Dh] → Output: [Tq, Dh]
    pub fn attention(&self, k: &Tensor, v: &Tensor, scale: f32, causal: bool) -> Result<Tensor> {
        if self.shape.ndim() != 2 || k.shape.ndim() != 2 || v.shape.ndim() != 2 {
            return Err(MlxError::InvalidArgument(
                "attention requires 2D tensors [seq, head_dim]".into(),
            ));
        }
        let tq = self.shape.0[0];
        let dh = self.shape.0[1];
        if k.shape.0[1] != dh {
            return Err(MlxError::ShapeMismatch {
                expected: self.shape.0.clone(),
                got: k.shape.0.clone(),
            });
        }
        if v.shape.0[1] != dh || k.shape.0[0] != v.shape.0[0] {
            return Err(MlxError::ShapeMismatch {
                expected: k.shape.0.clone(),
                got: v.shape.0.clone(),
            });
        }
        Ok(self.lazy_op(
            OpKind::Attention { scale, causal },
            SmallVec::from_slice(&[self.node_id, k.node_id, v.node_id]),
            Shape::new(vec![tq, dh]),
            self.dtype,
        ))
    }

    /// Element-wise square root.
    pub fn sqrt(&self) -> Tensor {
        self.lazy_op(
            OpKind::Sqrt,
            SmallVec::from_slice(&[self.node_id]),
            self.shape.clone(),
            self.dtype,
        )
    }

    // ── Materialization ─────────────────────────────────────────────────

    /// Materialize the tensor — triggers evaluation of the computation graph.
    pub fn eval(&self) -> Result<()> {
        self.stream.eval(self.node_id)
    }

    /// Copy data out as Vec<f32>. Triggers evaluation if needed.
    pub fn to_vec_f32(&self) -> Result<Vec<f32>> {
        self.eval()?;
        self.stream
            .get_buffer(self.node_id)
            .ok_or_else(|| MlxError::InvalidArgument("buffer not found after eval".into()))
    }

    // ── Accessors ───────────────────────────────────────────────────────

    /// Get the tensor shape.
    pub fn shape(&self) -> &Shape {
        &self.shape
    }

    /// Get the tensor dtype.
    pub fn dtype(&self) -> DType {
        self.dtype
    }

    /// Get the tensor device.
    pub fn device(&self) -> &Device {
        &self.device
    }

    /// Number of elements.
    pub fn numel(&self) -> i64 {
        self.shape.numel()
    }

    /// Get the graph node ID.
    pub fn node_id(&self) -> NodeId {
        self.node_id
    }

    /// Get the stream this tensor belongs to.
    pub fn stream(&self) -> Arc<Stream> {
        Arc::clone(&self.stream)
    }

    /// Reconstruct a tensor handle from a node ID and metadata.
    ///
    /// Used by autograd to create handles for graph introspection.
    pub fn from_node_id(
        node_id: NodeId,
        shape: Shape,
        dtype: DType,
        device: Device,
        stream: Arc<Stream>,
    ) -> Self {
        Self {
            node_id,
            shape,
            dtype,
            device,
            stream,
        }
    }

    /// Broadcast this tensor to the target shape (numpy-style rules).
    pub fn broadcast_to(&self, target: &Shape) -> Result<Tensor> {
        if &self.shape == target {
            return Ok(self.clone());
        }
        // Validate broadcast compatibility: dimensions are compared from the right.
        let in_ndim = self.shape.ndim();
        let out_ndim = target.ndim();
        if in_ndim > out_ndim {
            return Err(MlxError::InvalidArgument(format!(
                "cannot broadcast shape {} to {}",
                self.shape, target
            )));
        }
        let pad = out_ndim - in_ndim;
        for i in 0..in_ndim {
            let in_dim = self.shape.0[i];
            let out_dim = target.0[pad + i];
            if in_dim != 1 && in_dim != out_dim {
                return Err(MlxError::InvalidArgument(format!(
                    "cannot broadcast shape {} to {}",
                    self.shape, target
                )));
            }
        }
        Ok(self.lazy_op(
            OpKind::Broadcast {
                target_shape: target.clone(),
            },
            SmallVec::from_slice(&[self.node_id]),
            target.clone(),
            self.dtype,
        ))
    }
}

impl std::ops::Add for &Tensor {
    type Output = Result<Tensor>;
    fn add(self, rhs: &Tensor) -> Self::Output {
        self.add(rhs)
    }
}

impl std::ops::Sub for &Tensor {
    type Output = Result<Tensor>;
    fn sub(self, rhs: &Tensor) -> Self::Output {
        Tensor::sub(self, rhs)
    }
}

impl std::ops::Mul for &Tensor {
    type Output = Result<Tensor>;
    fn mul(self, rhs: &Tensor) -> Self::Output {
        Tensor::mul(self, rhs)
    }
}

impl std::ops::Neg for &Tensor {
    type Output = Tensor;
    fn neg(self) -> Self::Output {
        Tensor::neg(self)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn cpu() -> Device {
        Device::Cpu
    }

    #[test]
    fn test_zeros() {
        let t = Tensor::zeros(&Shape::new(vec![2, 3]), DType::F32, &cpu()).unwrap();
        assert_eq!(t.to_vec_f32().unwrap(), vec![0.0; 6]);
        assert_eq!(t.shape(), &Shape::new(vec![2, 3]));
    }

    #[test]
    fn test_ones() {
        let t = Tensor::ones(&Shape::new(vec![3]), DType::F32, &cpu()).unwrap();
        assert_eq!(t.to_vec_f32().unwrap(), vec![1.0; 3]);
    }

    #[test]
    fn test_from_f32() {
        let t = Tensor::from_f32(&[1.0, 2.0, 3.0, 4.0], &Shape::new(vec![2, 2]), &cpu()).unwrap();
        assert_eq!(t.to_vec_f32().unwrap(), vec![1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn test_from_f32_shape_mismatch() {
        let r = Tensor::from_f32(&[1.0, 2.0], &Shape::new(vec![3]), &cpu());
        assert!(r.is_err());
    }

    #[test]
    fn test_add() {
        let a = Tensor::from_f32(&[1.0, 2.0, 3.0], &Shape::new(vec![3]), &cpu()).unwrap();
        let b = Tensor::from_f32(&[4.0, 5.0, 6.0], &Shape::new(vec![3]), &cpu()).unwrap();
        let c = a.add(&b).unwrap();
        assert_eq!(c.to_vec_f32().unwrap(), vec![5.0, 7.0, 9.0]);
    }

    #[test]
    fn test_sub() {
        let a = Tensor::from_f32(&[5.0, 7.0, 9.0], &Shape::new(vec![3]), &cpu()).unwrap();
        let b = Tensor::from_f32(&[1.0, 2.0, 3.0], &Shape::new(vec![3]), &cpu()).unwrap();
        let c = a.sub(&b).unwrap();
        assert_eq!(c.to_vec_f32().unwrap(), vec![4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_mul() {
        let a = Tensor::from_f32(&[2.0, 3.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let b = Tensor::from_f32(&[4.0, 5.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let c = a.mul(&b).unwrap();
        assert_eq!(c.to_vec_f32().unwrap(), vec![8.0, 15.0]);
    }

    #[test]
    fn test_div() {
        let a = Tensor::from_f32(&[10.0, 9.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let b = Tensor::from_f32(&[2.0, 3.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let c = a.div(&b).unwrap();
        assert_eq!(c.to_vec_f32().unwrap(), vec![5.0, 3.0]);
    }

    #[test]
    fn test_neg() {
        let a = Tensor::from_f32(&[1.0, -2.0, 3.0], &Shape::new(vec![3]), &cpu()).unwrap();
        let b = a.neg();
        assert_eq!(b.to_vec_f32().unwrap(), vec![-1.0, 2.0, -3.0]);
    }

    #[test]
    fn test_matmul() {
        let a = Tensor::from_f32(&[1.0, 2.0, 3.0, 4.0], &Shape::new(vec![2, 2]), &cpu()).unwrap();
        let b = Tensor::from_f32(&[5.0, 6.0, 7.0, 8.0], &Shape::new(vec![2, 2]), &cpu()).unwrap();
        let c = a.matmul(&b).unwrap();
        assert_eq!(c.to_vec_f32().unwrap(), vec![19.0, 22.0, 43.0, 50.0]);
    }

    #[test]
    fn test_sum_axis() {
        let a = Tensor::from_f32(
            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            &Shape::new(vec![2, 3]),
            &cpu(),
        )
        .unwrap();
        let s0 = a.sum_axis(0).unwrap();
        assert_eq!(s0.to_vec_f32().unwrap(), vec![5.0, 7.0, 9.0]);
        let s1 = a.sum_axis(1).unwrap();
        assert_eq!(s1.to_vec_f32().unwrap(), vec![6.0, 15.0]);
    }

    #[test]
    fn test_sum_all() {
        let a = Tensor::from_f32(&[1.0, 2.0, 3.0], &Shape::new(vec![3]), &cpu()).unwrap();
        let s = a.sum_all().unwrap();
        assert_eq!(s.to_vec_f32().unwrap(), vec![6.0]);
    }

    #[test]
    fn test_softmax() {
        let a = Tensor::from_f32(&[1.0, 2.0, 3.0], &Shape::new(vec![3]), &cpu()).unwrap();
        let s = a.softmax(0).unwrap();
        let vals = s.to_vec_f32().unwrap();
        let sum: f32 = vals.iter().sum();
        assert!((sum - 1.0).abs() < 1e-6);
        assert!(vals[0] < vals[1]);
        assert!(vals[1] < vals[2]);
    }

    #[test]
    fn test_reshape() {
        let a = Tensor::from_f32(
            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            &Shape::new(vec![2, 3]),
            &cpu(),
        )
        .unwrap();
        let b = a.reshape(&Shape::new(vec![3, 2])).unwrap();
        assert_eq!(b.shape(), &Shape::new(vec![3, 2]));
        assert_eq!(b.to_vec_f32().unwrap(), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_transpose() {
        let a = Tensor::from_f32(
            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            &Shape::new(vec![2, 3]),
            &cpu(),
        )
        .unwrap();
        let b = a.transpose(None).unwrap();
        assert_eq!(b.shape(), &Shape::new(vec![3, 2]));
        // [[1,2,3],[4,5,6]] transposed = [[1,4],[2,5],[3,6]]
        assert_eq!(b.to_vec_f32().unwrap(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
    }

    #[test]
    fn test_operator_add() {
        let a = Tensor::from_f32(&[1.0, 2.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let b = Tensor::from_f32(&[3.0, 4.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let c = (&a + &b).unwrap();
        assert_eq!(c.to_vec_f32().unwrap(), vec![4.0, 6.0]);
    }

    #[test]
    fn test_operator_neg() {
        let a = Tensor::from_f32(&[1.0, -2.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let b = -&a;
        assert_eq!(b.to_vec_f32().unwrap(), vec![-1.0, 2.0]);
    }

    #[test]
    fn test_lazy_chain() {
        // Build a chain: (a + b) * c — nothing evaluated until to_vec_f32
        let a = Tensor::from_f32(&[1.0, 2.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let b = Tensor::from_f32(&[3.0, 4.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let c = Tensor::from_f32(&[2.0, 3.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let d = a.add(&b).unwrap().mul(&c).unwrap();
        // Only now does evaluation happen:
        assert_eq!(d.to_vec_f32().unwrap(), vec![8.0, 18.0]);
    }

    #[test]
    fn test_silu() {
        let a = Tensor::from_f32(&[0.0, 1.0], &Shape::new(vec![2]), &cpu()).unwrap();
        let b = a.silu();
        let vals = b.to_vec_f32().unwrap();
        assert!((vals[0]).abs() < 1e-6);
        assert!((vals[1] - 0.7311).abs() < 1e-3);
    }

    #[test]
    fn test_layer_norm() {
        let a = Tensor::from_f32(&[1.0, 2.0, 3.0], &Shape::new(vec![3]), &cpu()).unwrap();
        let b = a.layer_norm(1e-5);
        let vals = b.to_vec_f32().unwrap();
        let mean: f32 = vals.iter().sum::<f32>() / 3.0;
        assert!(mean.abs() < 1e-5);
    }

    #[test]
    fn test_reduce_zero_dim_bug() {
        let x = Tensor::from_f32(&[], &Shape::new(vec![2, 3, 0]), &cpu()).unwrap();
        let s = x.sum_axis(1).unwrap(); // Should return shape [2, 0]
        assert_eq!(s.shape(), &Shape::new(vec![2, 0]));
        let vals = s.to_vec_f32().unwrap();
        assert_eq!(vals.len(), 0);
    }

    #[test]
    fn test_softmax_zero_trailing_dim() {
        let x = Tensor::from_f32(&[], &Shape::new(vec![2, 3, 0]), &cpu()).unwrap();
        let s = x.softmax(1).unwrap();
        assert_eq!(s.shape(), &Shape::new(vec![2, 3, 0]));
        let vals = s.to_vec_f32().unwrap();
        assert_eq!(vals.len(), 0);
    }
}