grafos_tensor/

lib.rs

//! grafos-tensor -- Tensor operations backed by leased fabric memory.
//!
//! This crate provides [`FabricTensor`], an N-dimensional tensor whose data is
//! stored as contiguous row-major `f32` values backed by fabric memory acquired
//! through FBMU (Fabric Bootstrap Memory Unit) leases. Every constructor
//! validates that the fabric has sufficient capacity before allocating, and
//! **holds the lease for the lifetime of the tensor**.
//!
//! # Supported operations
//!
//! | Operation | Method | Constraints |
//! |-----------|--------|-------------|
//! | Matrix multiply | [`FabricTensor::matmul`] | Both 2-D; inner dims match |
//! | Elementwise add | [`FabricTensor::add`] | Same shape |
//! | Elementwise mul | [`FabricTensor::mul`] | Same shape |
//! | Scalar multiply | [`FabricTensor::scale`] | Any shape |
//! | ReLU | [`FabricTensor::relu`] | Any shape |
//! | Softmax | [`FabricTensor::softmax`] | `axis < ndim` |
//! | Subtract | [`FabricTensor::subtract`] | Same shape |
//! | Sum axis | [`FabricTensor::sum_axis`] | `axis < ndim` |
//! | Sigmoid | [`FabricTensor::sigmoid`] | Any shape |
//! | Natural log | [`FabricTensor::ln`] | Any shape |
//! | Clamp | [`FabricTensor::clip`] | Any shape |
//! | Transpose | [`FabricTensor::transpose`] | ndim >= 2 |
//! | Reshape | [`FabricTensor::reshape`] | Same total elements |
//!
//! Placement helpers:
//! - [`FabricTensor::to_gpu`] / [`FabricTensor::to_cpu`]
//! - [`FabricTensor::device`], [`FabricTensor::is_cpu`], [`FabricTensor::is_gpu`]
//!
//! With `gpu` feature enabled, operations on GPU-placed tensors dispatch
//! through the v1 `fabricbios_gpu_v1` session surface — specifically a
//! private `submit_signal_kernel` helper that opens a transient
//! [`grafos_std::gpu::GpuSession`] on the tensor's held `GpuLease`,
//! `module_load`s the kernel binary, `launch`es with `[1,1,1]` grid/block
//! dims, `sync`s, and lets the session/module drop (RAII). Numerical
//! outputs are still computed on the CPU path; GPU dispatch is
//! control-path signaling only under the current mock kernels. See
//! `docs/grafos-tensor-guide.md` for the full programming model.
//!
//! Operator overloading is provided for `&FabricTensor`: `+` (elementwise add),
//! `*` (elementwise mul), and `* f32` (scalar mul).
//!
//! # Quick start
//!
//! ```rust
//! use grafos_tensor::FabricTensor;
//!
//! # grafos_std::host::reset_mock();
//! # grafos_std::host::mock_set_fbmu_arena_size(65536);
//! let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
//! let b = FabricTensor::from_slice(&[3, 2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).unwrap();
//! let c = a.matmul(&b).unwrap();
//! assert_eq!(c.shape(), &[2, 2]);
//! assert_eq!(c.get(&[0, 0]).unwrap(), 58.0);
//! ```
//!
//! # Operator overloading
//!
//! ```rust
//! use grafos_tensor::FabricTensor;
//!
//! # grafos_std::host::reset_mock();
//! # grafos_std::host::mock_set_fbmu_arena_size(65536);
//! let a = FabricTensor::from_slice(&[3], &[1.0, 2.0, 3.0]).unwrap();
//! let b = FabricTensor::from_slice(&[3], &[4.0, 5.0, 6.0]).unwrap();
//!
//! let sum = (&a + &b).unwrap();        // elementwise add
//! let prod = (&a * &b).unwrap();       // elementwise mul
//! let scaled = (&a * 2.0f32).unwrap(); // scalar mul
//! ```
//!
//! # Testing
//!
//! On native targets, initialize the mock FBMU backend before creating tensors:
//!
//! ```rust
//! grafos_std::host::reset_mock();
//! grafos_std::host::mock_set_fbmu_arena_size(1 << 20); // 1 MiB
//! ```

#![cfg_attr(not(feature = "std"), no_std)]

extern crate alloc;

#[cfg(feature = "gpu")]
use alloc::borrow::Cow;
use alloc::vec;
use alloc::vec::Vec;
use core::ops::{Add, Mul, Sub};

use grafos_std::error::FabricError;
#[cfg(feature = "gpu")]
use grafos_std::gpu::{GpuBuilder, GpuLease, GpuSession};
use grafos_std::mem::{MemBuilder, MemLease};
#[cfg(all(feature = "gpu", feature = "std"))]
use std::{env, fs, path::PathBuf};

/// Result alias using [`FabricError`].
pub type Result<T> = core::result::Result<T, FabricError>;

/// N-dimensional tensor backed by a fabric memory lease.
///
/// Tensor data is stored as a contiguous row-major `f32` array. Each
/// constructor acquires a [`MemLease`] to validate
/// that the fabric has sufficient capacity. The [`Shape`] metadata tracks
/// dimensions and precomputed strides for multi-dimensional indexing.
///
/// # Storage layout
///
/// Data is laid out in row-major (C) order: the last dimension varies fastest
/// in memory. For a `[2, 3]` matrix `[[a, b, c], [d, e, f]]`, the flat
/// representation is `[a, b, c, d, e, f]` with strides `[3, 1]`.
///
/// # Examples
///
/// ```rust
/// use grafos_tensor::FabricTensor;
///
/// # grafos_std::host::reset_mock();
/// # grafos_std::host::mock_set_fbmu_arena_size(65536);
/// // Create a 2x3 matrix and access elements
/// let t = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
/// assert_eq!(t.get(&[1, 2]).unwrap(), 6.0);
/// assert_eq!(t.shape(), &[2, 3]);
/// assert_eq!(t.strides(), &[3, 1]);
/// ```
pub struct FabricTensor {
    storage: TensorStorage,
    shape: Shape,
    data: Vec<f32>,
}

/// Placement of tensor data in the fabric.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Device {
    /// Tensor is placed on CPU memory (FBMU-backed lease).
    Cpu,
    /// Tensor is placed on GPU memory (GPU lease id).
    Gpu(u128),
}

enum TensorStorage {
    Cpu(MemLease),
    #[cfg(feature = "gpu")]
    Gpu {
        gpu_lease: GpuLease,
        staging_lease: MemLease,
    },
}

/// Shape metadata for an N-dimensional tensor.
///
/// Stores dimension sizes and precomputed row-major strides. Strides follow
/// the standard row-major formula: `strides[n-1] = 1` and
/// `strides[i] = strides[i+1] * dims[i+1]` for preceding dimensions.
///
/// # Examples
///
/// ```rust
/// use grafos_tensor::Shape;
///
/// let s = Shape::new(&[2, 3, 4]);
/// assert_eq!(s.dims(), &[2, 3, 4]);
/// assert_eq!(s.strides(), &[12, 4, 1]);
/// assert_eq!(s.numel(), 24);
/// assert_eq!(s.ndim(), 3);
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Shape {
    dims: Vec<usize>,
    strides: Vec<usize>,
}

impl Shape {
    /// Create a new shape from dimension sizes.
    ///
    /// Strides are computed in row-major (C) order: the last dimension
    /// has stride 1, and each preceding dimension's stride is the product
    /// of all following dimension sizes.
    ///
    /// An empty `dims` slice creates a scalar (0-D) shape with `numel() == 1`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::Shape;
    ///
    /// let s = Shape::new(&[2, 3]);
    /// assert_eq!(s.strides(), &[3, 1]);
    ///
    /// let scalar = Shape::new(&[]);
    /// assert_eq!(scalar.numel(), 1);
    /// ```
    pub fn new(dims: &[usize]) -> Self {
        let strides = Self::compute_strides(dims);
        Shape {
            dims: dims.to_vec(),
            strides,
        }
    }

    fn compute_strides(dims: &[usize]) -> Vec<usize> {
        if dims.is_empty() {
            return Vec::new();
        }
        let mut strides = vec![1usize; dims.len()];
        for i in (0..dims.len() - 1).rev() {
            strides[i] = strides[i + 1] * dims[i + 1];
        }
        strides
    }

    /// Total number of elements in the tensor.
    pub fn numel(&self) -> usize {
        self.dims.iter().product()
    }

    /// Number of dimensions.
    pub fn ndim(&self) -> usize {
        self.dims.len()
    }

    /// Dimension sizes.
    pub fn dims(&self) -> &[usize] {
        &self.dims
    }

    /// Row-major strides.
    pub fn strides(&self) -> &[usize] {
        &self.strides
    }

    /// Convert multi-dimensional indices to a flat offset.
    ///
    /// Returns `None` if any index is out of bounds.
    fn flat_index(&self, indices: &[usize]) -> Option<usize> {
        if indices.len() != self.dims.len() {
            return None;
        }
        let mut offset = 0;
        for (i, &idx) in indices.iter().enumerate() {
            if idx >= self.dims[i] {
                return None;
            }
            offset += idx * self.strides[i];
        }
        Some(offset)
    }
}

/// Execute a single signal-only GPU kernel dispatch via the v1
/// `GpuSession` surface.
///
/// grafos-tensor's GPU dispatch is control-path signaling only: kernel
/// binaries default to the `b"grafos.tensor.mock"` stub and numerical
/// results are computed by the CPU path (see
/// `docs/grafos-tensor-guide.md`). This helper encapsulates the canonical
/// v1 sequence — `module_load` → `launch` → `sync` — with the
/// `GpuModule` and `GpuSession` dropping at function exit (RAII unload +
/// release of the borrowed lease).
///
/// Grid and block default to `[1, 1, 1]` because the mock kernel does
/// not read them. Real kernels would require caller-provided launch
/// geometry, which is deferred until grafos-tensor ships real kernels.
#[cfg(feature = "gpu")]
fn submit_signal_kernel(
    gpu_lease: &GpuLease,
    op_name: &str,
    binary: &[u8],
    args: &[u8],
    arg_sizes: &[u32],
) -> Result<()> {
    let mut sess = GpuSession::new(gpu_lease);
    let module = sess.module_load(binary)?;
    sess.launch(&module, op_name, [1, 1, 1], [1, 1, 1], args, arg_sizes)?;
    sess.sync()?;
    Ok(())
}

impl FabricTensor {
    fn cpu_lease(&self) -> &MemLease {
        match &self.storage {
            TensorStorage::Cpu(lease) => lease,
            #[cfg(feature = "gpu")]
            TensorStorage::Gpu { staging_lease, .. } => staging_lease,
        }
    }

    #[cfg(feature = "gpu")]
    fn gpu_lease(&self) -> Option<&GpuLease> {
        match &self.storage {
            TensorStorage::Cpu(_) => None,
            TensorStorage::Gpu { gpu_lease, .. } => Some(gpu_lease),
        }
    }

    /// Access the underlying fabric memory lease held by this tensor.
    ///
    /// The current implementation stores element data in a local `Vec<f32>`,
    /// but holds a [`MemLease`] as the capacity/liveness contract for the
    /// tensor's lifetime.
    pub fn lease(&self) -> &MemLease {
        self.cpu_lease()
    }

    /// Wrap an existing memory lease as a tensor.
    ///
    /// The tensor takes ownership of `lease` and uses it as its backing
    /// memory contract. `data` is initialized to zeros.
    pub fn from_mem_lease(shape: &[usize], lease: MemLease) -> Self {
        let shape = Shape::new(shape);
        let numel = shape.numel();
        FabricTensor {
            storage: TensorStorage::Cpu(lease),
            shape,
            data: vec![0.0; numel],
        }
    }

    /// Create a tensor filled with zeros.
    ///
    /// Acquires a [`MemLease`] large enough for
    /// `numel * 4` bytes and initializes all elements to `0.0`.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if the fabric cannot provide
    /// sufficient memory, or [`FabricError::Disconnected`] if the FBMU
    /// handshake fails.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let t = FabricTensor::zeros(&[3, 4]).unwrap();
    /// assert_eq!(t.shape(), &[3, 4]);
    /// assert_eq!(t.numel(), 12);
    /// assert_eq!(t.get(&[0, 0]).unwrap(), 0.0);
    /// ```
    pub fn zeros(shape: &[usize]) -> Result<Self> {
        let s = Shape::new(shape);
        let numel = s.numel();
        let byte_size = numel * core::mem::size_of::<f32>();
        // Acquire and hold a lease to prove (and reserve) fabric capacity for
        // the lifetime of this tensor. The current implementation stores the
        // actual element data locally, but the lease is still the liveness and
        // capacity contract.
        let lease = MemBuilder::new().min_bytes(byte_size as u64).acquire()?;
        Ok(FabricTensor {
            storage: TensorStorage::Cpu(lease),
            shape: s,
            data: vec![0.0f32; numel],
        })
    }

    /// Create a tensor from a slice of data.
    ///
    /// The length of `data` must exactly equal the product of `shape`
    /// dimensions. Data is copied into the tensor in row-major order.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if `data.len()` does not
    /// match the shape's element count, or if the fabric cannot provide
    /// sufficient memory.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let t = FabricTensor::from_slice(&[2, 2], &[1.0, 2.0, 3.0, 4.0]).unwrap();
    /// assert_eq!(t.get(&[0, 1]).unwrap(), 2.0);
    /// assert_eq!(t.get(&[1, 0]).unwrap(), 3.0);
    /// ```
    pub fn from_slice(shape: &[usize], data: &[f32]) -> Result<Self> {
        let s = Shape::new(shape);
        if data.len() != s.numel() {
            return Err(FabricError::CapacityExceeded);
        }
        let byte_size = std::mem::size_of_val(data);
        let lease = MemBuilder::new().min_bytes(byte_size as u64).acquire()?;
        Ok(FabricTensor {
            storage: TensorStorage::Cpu(lease),
            shape: s,
            data: data.to_vec(),
        })
    }

    /// Return the current placement device for this tensor.
    pub fn device(&self) -> Device {
        #[cfg(feature = "gpu")]
        {
            if let Some(gpu) = self.gpu_lease() {
                return Device::Gpu(gpu.lease_id());
            }
        }
        Device::Cpu
    }

    /// Returns `true` if the tensor is placed on CPU memory.
    pub fn is_cpu(&self) -> bool {
        matches!(self.device(), Device::Cpu)
    }

    /// Returns `true` if the tensor is placed on GPU memory.
    pub fn is_gpu(&self) -> bool {
        matches!(self.device(), Device::Gpu(_))
    }

    /// Move/copy this tensor to GPU placement.
    ///
    /// On builds without the `gpu` feature, returns [`FabricError::Unsupported`].
    pub fn to_gpu(&self) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        {
            if self.is_gpu() {
                return FabricTensor::from_slice(self.shape(), self.as_slice()).and_then(
                    |mut t| {
                        let gpu = GpuBuilder::new()
                            .min_vram((self.numel() * core::mem::size_of::<f32>()) as u64)
                            .acquire()?;
                        let staging = MemBuilder::new()
                            .min_bytes((self.numel() * core::mem::size_of::<f32>()) as u64)
                            .acquire()?;
                        let args = self.encode_f32_le();
                        let arg_sizes = [args.len() as u32];
                        submit_signal_kernel(
                            &gpu,
                            "tensor_upload",
                            b"grafos.tensor.mock",
                            &args,
                            &arg_sizes,
                        )?;
                        t.storage = TensorStorage::Gpu {
                            gpu_lease: gpu,
                            staging_lease: staging,
                        };
                        Ok(t)
                    },
                );
            }

            let bytes = (self.numel() * core::mem::size_of::<f32>()) as u64;
            let gpu_lease = GpuBuilder::new().min_vram(bytes).acquire()?;
            let staging_lease = MemBuilder::new().min_bytes(bytes).acquire()?;
            let args = self.encode_f32_le();
            let arg_sizes = [args.len() as u32];
            submit_signal_kernel(
                &gpu_lease,
                "tensor_upload",
                b"grafos.tensor.mock",
                &args,
                &arg_sizes,
            )?;

            return Ok(FabricTensor {
                storage: TensorStorage::Gpu {
                    gpu_lease,
                    staging_lease,
                },
                shape: self.shape.clone(),
                data: self.data.clone(),
            });
        }

        #[cfg(not(feature = "gpu"))]
        {
            Err(FabricError::Unsupported)
        }
    }

    /// Move/copy this tensor to CPU placement.
    pub fn to_cpu(&self) -> Result<FabricTensor> {
        if self.is_cpu() {
            return FabricTensor::from_slice(self.shape(), self.as_slice());
        }

        #[cfg(feature = "gpu")]
        {
            if let Some(gpu) = self.gpu_lease() {
                let args = self.encode_f32_le();
                let arg_sizes = [args.len() as u32];
                submit_signal_kernel(
                    gpu,
                    "tensor_download",
                    b"grafos.tensor.mock",
                    &args,
                    &arg_sizes,
                )?;
            }
        }
        FabricTensor::from_slice(self.shape(), self.as_slice())
    }

    /// Get the element at the given multi-dimensional indices.
    ///
    /// The length of `indices` must equal [`ndim()`](Self::ndim), and each
    /// index must be within its dimension's range.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if any index is out of
    /// bounds or if the wrong number of indices is provided.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let t = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    /// assert_eq!(t.get(&[0, 2]).unwrap(), 3.0);
    /// assert_eq!(t.get(&[1, 1]).unwrap(), 5.0);
    /// assert!(t.get(&[2, 0]).is_err()); // out of bounds
    /// ```
    pub fn get(&self, indices: &[usize]) -> Result<f32> {
        let offset = self
            .shape
            .flat_index(indices)
            .ok_or(FabricError::CapacityExceeded)?;
        Ok(self.data[offset])
    }

    /// Set the element at the given multi-dimensional indices.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if any index is out of
    /// bounds or if the wrong number of indices is provided.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let mut t = FabricTensor::zeros(&[2, 2]).unwrap();
    /// t.set(&[1, 0], 42.0).unwrap();
    /// assert_eq!(t.get(&[1, 0]).unwrap(), 42.0);
    /// ```
    pub fn set(&mut self, indices: &[usize], value: f32) -> Result<()> {
        let offset = self
            .shape
            .flat_index(indices)
            .ok_or(FabricError::CapacityExceeded)?;
        self.data[offset] = value;
        Ok(())
    }

    /// Dimension sizes as a slice (e.g. `&[2, 3]` for a 2x3 matrix).
    pub fn shape(&self) -> &[usize] {
        self.shape.dims()
    }

    /// Number of dimensions (0 for scalar, 1 for vector, 2 for matrix, etc.).
    pub fn ndim(&self) -> usize {
        self.shape.ndim()
    }

    /// Total number of elements (product of all dimension sizes).
    pub fn numel(&self) -> usize {
        self.shape.numel()
    }

    /// Row-major strides (e.g. `&[3, 1]` for a 2x3 matrix).
    pub fn strides(&self) -> &[usize] {
        self.shape.strides()
    }

    /// Access raw data as a contiguous `f32` slice in row-major order.
    pub fn as_slice(&self) -> &[f32] {
        &self.data
    }

    #[cfg(feature = "gpu")]
    fn encode_f32_le(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(self.data.len() * core::mem::size_of::<f32>());
        for value in &self.data {
            bytes.extend_from_slice(&value.to_le_bytes());
        }
        bytes
    }

    #[cfg(feature = "gpu")]
    fn push_u16_le(buf: &mut Vec<u8>, v: u16) {
        buf.extend_from_slice(&v.to_le_bytes());
    }

    #[cfg(feature = "gpu")]
    fn push_u32_le(buf: &mut Vec<u8>, v: u32) {
        buf.extend_from_slice(&v.to_le_bytes());
    }

    #[cfg(feature = "gpu")]
    fn push_u64_le(buf: &mut Vec<u8>, v: u64) {
        buf.extend_from_slice(&v.to_le_bytes());
    }

    #[cfg(feature = "gpu")]
    fn encode_tensor_descriptor(&self, out: &mut Vec<u8>) -> Result<()> {
        let ndim: u16 = self
            .ndim()
            .try_into()
            .map_err(|_| FabricError::CapacityExceeded)?;
        Self::push_u16_le(out, ndim);

        let numel: u64 = self
            .numel()
            .try_into()
            .map_err(|_| FabricError::CapacityExceeded)?;
        Self::push_u64_le(out, numel);
        Self::push_u64_le(out, 0); // offset

        for &dim in self.shape() {
            let dim_u32: u32 = dim.try_into().map_err(|_| FabricError::CapacityExceeded)?;
            Self::push_u32_le(out, dim_u32);
        }
        for &stride in self.strides() {
            let stride_u32: u32 = stride
                .try_into()
                .map_err(|_| FabricError::CapacityExceeded)?;
            Self::push_u32_le(out, stride_u32);
        }
        Ok(())
    }

    #[cfg(feature = "gpu")]
    fn kernel_binary_for(op_name: &str) -> Cow<'static, [u8]> {
        #[cfg(all(feature = "gpu", feature = "std"))]
        {
            if let Some(file) = env::var_os("GRAFOS_TENSOR_HSACO") {
                let path = PathBuf::from(file);
                if let Ok(binary) = fs::read(path) {
                    if !binary.is_empty() {
                        return Cow::Owned(binary);
                    }
                }
            }

            if let Some(dir) = env::var_os("GRAFOS_TENSOR_KERNEL_DIR") {
                let base = PathBuf::from(dir);
                let candidates = [
                    base.join(format!("{op_name}.hsaco")),
                    base.join("tensor_ops_gfx1032.hsaco"),
                    base.join("tensor_ops.hsaco"),
                ];
                for candidate in candidates {
                    if let Ok(binary) = fs::read(&candidate) {
                        if !binary.is_empty() {
                            return Cow::Owned(binary);
                        }
                    }
                }
            }
        }
        Cow::Borrowed(b"grafos.tensor.mock")
    }

    #[cfg(feature = "gpu")]
    fn pack_gpu_unary_args(&self, op_name: &str) -> Result<Vec<u8>> {
        let mut args = Vec::new();
        args.extend_from_slice(b"GTA0");
        args.push(1); // version
        args.push(1); // tensor count
        Self::push_u16_le(&mut args, op_name.len() as u16);
        args.extend_from_slice(op_name.as_bytes());
        self.encode_tensor_descriptor(&mut args)?;
        let data = self.encode_f32_le();
        Self::push_u64_le(&mut args, data.len() as u64);
        args.extend_from_slice(&data);
        Ok(args)
    }

    #[cfg(feature = "gpu")]
    fn pack_gpu_binary_args(&self, other: &FabricTensor, op_name: &str) -> Result<Vec<u8>> {
        let mut args = Vec::new();
        args.extend_from_slice(b"GTA0");
        args.push(1); // version
        args.push(2); // tensor count
        Self::push_u16_le(&mut args, op_name.len() as u16);
        args.extend_from_slice(op_name.as_bytes());

        self.encode_tensor_descriptor(&mut args)?;
        let lhs = self.encode_f32_le();
        Self::push_u64_le(&mut args, lhs.len() as u64);
        args.extend_from_slice(&lhs);

        other.encode_tensor_descriptor(&mut args)?;
        let rhs = other.encode_f32_le();
        Self::push_u64_le(&mut args, rhs.len() as u64);
        args.extend_from_slice(&rhs);
        Ok(args)
    }

    #[cfg(feature = "gpu")]
    fn dispatch_gpu_unary(&self, op_name: &str) -> Result<()> {
        let Some(gpu) = self.gpu_lease() else {
            return Ok(());
        };
        let args = self.pack_gpu_unary_args(op_name)?;
        let binary = Self::kernel_binary_for(op_name);
        let arg_sizes = [args.len() as u32];
        submit_signal_kernel(gpu, op_name, binary.as_ref(), &args, &arg_sizes)?;
        Ok(())
    }

    #[cfg(feature = "gpu")]
    fn dispatch_gpu_binary(&self, other: &FabricTensor, op_name: &str) -> Result<()> {
        let Some(gpu) = self.gpu_lease() else {
            return Ok(());
        };
        let args = self.pack_gpu_binary_args(other, op_name)?;
        let binary = Self::kernel_binary_for(op_name);
        let arg_sizes = [args.len() as u32];
        submit_signal_kernel(gpu, op_name, binary.as_ref(), &args, &arg_sizes)?;
        Ok(())
    }

    fn cpu_matmul(&self, other: &FabricTensor) -> Result<FabricTensor> {
        if self.ndim() != 2 || other.ndim() != 2 {
            return Err(FabricError::CapacityExceeded);
        }
        let m = self.shape.dims[0];
        let k = self.shape.dims[1];
        let k2 = other.shape.dims[0];
        let n = other.shape.dims[1];
        if k != k2 {
            return Err(FabricError::CapacityExceeded);
        }
        let mut result_data = vec![0.0f32; m * n];
        for i in 0..m {
            for j in 0..n {
                let mut sum = 0.0f32;
                for p in 0..k {
                    sum += self.data[i * k + p] * other.data[p * n + j];
                }
                result_data[i * n + j] = sum;
            }
        }
        FabricTensor::from_slice(&[m, n], &result_data)
    }

    fn cpu_add(&self, other: &FabricTensor) -> Result<FabricTensor> {
        if self.shape.dims != other.shape.dims {
            return Err(FabricError::CapacityExceeded);
        }
        let data: Vec<f32> = self
            .data
            .iter()
            .zip(other.data.iter())
            .map(|(a, b)| a + b)
            .collect();
        FabricTensor::from_slice(self.shape.dims(), &data)
    }

    fn cpu_mul(&self, other: &FabricTensor) -> Result<FabricTensor> {
        if self.shape.dims != other.shape.dims {
            return Err(FabricError::CapacityExceeded);
        }
        let data: Vec<f32> = self
            .data
            .iter()
            .zip(other.data.iter())
            .map(|(a, b)| a * b)
            .collect();
        FabricTensor::from_slice(self.shape.dims(), &data)
    }

    /// Matrix multiplication (2-D only).
    ///
    /// Computes `C = self @ other` using a triple-loop algorithm. Both tensors
    /// must be 2-D. If `self` has shape `[M, K]` and `other` has shape
    /// `[K, N]`, the result has shape `[M, N]`.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if either tensor is not 2-D,
    /// if the inner dimensions don't match, or if the fabric cannot provide
    /// memory for the result.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    /// let b = FabricTensor::from_slice(&[3, 2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).unwrap();
    /// let c = a.matmul(&b).unwrap();
    /// assert_eq!(c.shape(), &[2, 2]);
    /// // C[0,0] = 1*7 + 2*9 + 3*11 = 58
    /// assert_eq!(c.get(&[0, 0]).unwrap(), 58.0);
    /// ```
    pub fn matmul(&self, other: &FabricTensor) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() && other.is_gpu() {
            self.dispatch_gpu_binary(other, "tensor_matmul")?;
            let cpu_result = self.cpu_matmul(other)?;
            return cpu_result.to_gpu();
        }
        self.cpu_matmul(other)
    }

    /// Elementwise addition.
    ///
    /// Both tensors must have the same shape. Returns a new tensor where
    /// each element is the sum of the corresponding elements.
    ///
    /// Also available via the `+` operator: `(&a + &b).unwrap()`.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if shapes don't match.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[2, 2], &[1.0, 2.0, 3.0, 4.0]).unwrap();
    /// let b = FabricTensor::from_slice(&[2, 2], &[10.0, 20.0, 30.0, 40.0]).unwrap();
    /// let c = a.add(&b).unwrap();
    /// assert_eq!(c.as_slice(), &[11.0, 22.0, 33.0, 44.0]);
    /// ```
    pub fn add(&self, other: &FabricTensor) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() && other.is_gpu() {
            self.dispatch_gpu_binary(other, "tensor_add")?;
            let cpu_result = self.cpu_add(other)?;
            return cpu_result.to_gpu();
        }
        self.cpu_add(other)
    }

    /// Elementwise multiplication (Hadamard product).
    ///
    /// Both tensors must have the same shape. Returns a new tensor where
    /// each element is the product of the corresponding elements.
    ///
    /// Also available via the `*` operator: `(&a * &b).unwrap()`.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if shapes don't match.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[3], &[2.0, 3.0, 4.0]).unwrap();
    /// let b = FabricTensor::from_slice(&[3], &[5.0, 6.0, 7.0]).unwrap();
    /// let c = a.mul(&b).unwrap();
    /// assert_eq!(c.as_slice(), &[10.0, 18.0, 28.0]);
    /// ```
    pub fn mul(&self, other: &FabricTensor) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() && other.is_gpu() {
            self.dispatch_gpu_binary(other, "tensor_mul")?;
            let cpu_result = self.cpu_mul(other)?;
            return cpu_result.to_gpu();
        }
        self.cpu_mul(other)
    }

    /// Scalar multiplication.
    ///
    /// Multiplies every element by `scalar`. Works on tensors of any shape.
    ///
    /// Also available via the `*` operator with `f32`: `(&a * 2.0f32).unwrap()`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[2, 2], &[1.0, 2.0, 3.0, 4.0]).unwrap();
    /// let b = a.scale(3.0).unwrap();
    /// assert_eq!(b.as_slice(), &[3.0, 6.0, 9.0, 12.0]);
    /// ```
    pub fn scale(&self, scalar: f32) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_scale")?;
            let data: Vec<f32> = self.data.iter().map(|&x| x * scalar).collect();
            let cpu_result = FabricTensor::from_slice(self.shape.dims(), &data)?;
            return cpu_result.to_gpu();
        }
        let data: Vec<f32> = self.data.iter().map(|&x| x * scalar).collect();
        FabricTensor::from_slice(self.shape.dims(), &data)
    }

    /// ReLU activation: `max(0, x)` elementwise.
    ///
    /// Returns a new tensor where negative values are clamped to zero.
    /// Commonly used as an activation function in neural network layers.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[4], &[-2.0, -0.5, 0.0, 3.0]).unwrap();
    /// let b = a.relu().unwrap();
    /// assert_eq!(b.as_slice(), &[0.0, 0.0, 0.0, 3.0]);
    /// ```
    pub fn relu(&self) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_relu")?;
            let data: Vec<f32> = self
                .data
                .iter()
                .map(|&x| if x > 0.0 { x } else { 0.0 })
                .collect();
            let cpu_result = FabricTensor::from_slice(self.shape.dims(), &data)?;
            return cpu_result.to_gpu();
        }
        let data: Vec<f32> = self
            .data
            .iter()
            .map(|&x| if x > 0.0 { x } else { 0.0 })
            .collect();
        FabricTensor::from_slice(self.shape.dims(), &data)
    }

    /// Softmax along the specified axis.
    ///
    /// For each slice along `axis`, computes `exp(x - max) / sum(exp(x - max))`.
    /// The max-subtraction provides numerical stability against overflow.
    ///
    /// For a `[M, N]` matrix: `softmax(1)` normalizes each row (sums to 1),
    /// `softmax(0)` normalizes each column.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if `axis >= ndim()`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[1, 4], &[1.0, 2.0, 3.0, 4.0]).unwrap();
    /// let probs = a.softmax(1).unwrap();
    /// let sum: f32 = probs.as_slice().iter().sum();
    /// assert!((sum - 1.0).abs() < 1e-6);
    /// ```
    pub fn softmax(&self, axis: usize) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_softmax")?;
            let cpu_result = self.softmax_cpu(axis)?;
            return cpu_result.to_gpu();
        }
        self.softmax_cpu(axis)
    }

    /// Transpose: swap the last two dimensions.
    ///
    /// For a 2-D `[M, N]` matrix, produces `[N, M]`. For higher-rank tensors
    /// (e.g. `[B, M, N]`), produces `[B, N, M]` by batching over leading
    /// dimensions.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if the tensor has fewer than
    /// 2 dimensions.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    /// let b = a.transpose().unwrap();
    /// assert_eq!(b.shape(), &[3, 2]);
    /// assert_eq!(b.get(&[0, 1]).unwrap(), 4.0); // was a[1, 0]
    /// ```
    pub fn transpose(&self) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_transpose")?;
            let cpu_result = self.transpose_cpu()?;
            return cpu_result.to_gpu();
        }
        self.transpose_cpu()
    }

    /// Reshape to a new shape.
    ///
    /// The total number of elements must remain the same. Data is copied
    /// to a new tensor backed by a new lease; the underlying flat storage
    /// order does not change.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if the new shape's element
    /// count differs from the current tensor's.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    /// let b = a.reshape(&[3, 2]).unwrap();
    /// assert_eq!(b.shape(), &[3, 2]);
    /// assert_eq!(b.as_slice(), a.as_slice()); // same data order
    /// ```
    pub fn reshape(&self, new_shape: &[usize]) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_reshape")?;
            let cpu_result = self.reshape_cpu(new_shape)?;
            return cpu_result.to_gpu();
        }
        self.reshape_cpu(new_shape)
    }

    fn softmax_cpu(&self, axis: usize) -> Result<FabricTensor> {
        if axis >= self.ndim() {
            return Err(FabricError::CapacityExceeded);
        }
        let dims = &self.shape.dims;
        let strides = &self.shape.strides;
        let mut result = self.data.clone();

        let axis_len = dims[axis];
        let axis_stride = strides[axis];
        let outer_size: usize = dims[..axis].iter().product();
        let inner_size: usize = dims[axis + 1..].iter().product();

        for outer in 0..outer_size {
            for inner in 0..inner_size {
                let base = outer * (axis_len * inner_size) + inner;
                let mut max_val = f32::NEG_INFINITY;
                for a in 0..axis_len {
                    let idx = base + a * axis_stride;
                    if self.data[idx] > max_val {
                        max_val = self.data[idx];
                    }
                }

                let mut sum = 0.0f32;
                for a in 0..axis_len {
                    let idx = base + a * axis_stride;
                    let exp_val = (self.data[idx] - max_val).exp();
                    result[idx] = exp_val;
                    sum += exp_val;
                }

                for a in 0..axis_len {
                    let idx = base + a * axis_stride;
                    result[idx] /= sum;
                }
            }
        }

        FabricTensor::from_slice(dims, &result)
    }

    fn transpose_cpu(&self) -> Result<FabricTensor> {
        if self.ndim() < 2 {
            return Err(FabricError::CapacityExceeded);
        }
        let dims = &self.shape.dims;
        let ndim = dims.len();

        let mut new_dims = dims.to_vec();
        new_dims.swap(ndim - 2, ndim - 1);

        let rows = dims[ndim - 2];
        let cols = dims[ndim - 1];
        let batch_size: usize = dims[..ndim - 2].iter().product();
        let matrix_size = rows * cols;

        let mut result = vec![0.0f32; self.data.len()];
        for b in 0..batch_size {
            let src_base = b * matrix_size;
            let dst_base = b * matrix_size;
            for r in 0..rows {
                for c in 0..cols {
                    result[dst_base + c * rows + r] = self.data[src_base + r * cols + c];
                }
            }
        }

        FabricTensor::from_slice(&new_dims, &result)
    }

    fn reshape_cpu(&self, new_shape: &[usize]) -> Result<FabricTensor> {
        let new_numel: usize = new_shape.iter().product();
        if new_numel != self.numel() {
            return Err(FabricError::CapacityExceeded);
        }
        FabricTensor::from_slice(new_shape, &self.data)
    }

    /// Elementwise subtraction: `self - other`.
    ///
    /// Both tensors must have the same shape.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[3], &[5.0, 3.0, 1.0]).unwrap();
    /// let b = FabricTensor::from_slice(&[3], &[1.0, 2.0, 3.0]).unwrap();
    /// let c = a.subtract(&b).unwrap();
    /// assert_eq!(c.as_slice(), &[4.0, 1.0, -2.0]);
    /// ```
    pub fn subtract(&self, other: &FabricTensor) -> Result<FabricTensor> {
        if self.shape() != other.shape() {
            return Err(FabricError::CapacityExceeded);
        }
        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            let neg = other.scale(-1.0)?;
            return self.add(&neg);
        }
        let data: Vec<f32> = self
            .data
            .iter()
            .zip(other.data.iter())
            .map(|(&a, &b)| a - b)
            .collect();
        FabricTensor::from_slice(self.shape(), &data)
    }

    /// Sum along the specified axis, reducing that dimension.
    ///
    /// For a `[M, N]` matrix: `sum_axis(0)` produces `[N]` (column sums),
    /// `sum_axis(1)` produces `[M]` (row sums).
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if `axis >= ndim()`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    /// let row_sums = a.sum_axis(1).unwrap();
    /// assert_eq!(row_sums.shape(), &[2]);
    /// assert_eq!(row_sums.as_slice(), &[6.0, 15.0]);
    /// let col_sums = a.sum_axis(0).unwrap();
    /// assert_eq!(col_sums.shape(), &[3]);
    /// assert_eq!(col_sums.as_slice(), &[5.0, 7.0, 9.0]);
    /// ```
    pub fn sum_axis(&self, axis: usize) -> Result<FabricTensor> {
        if axis >= self.ndim() {
            return Err(FabricError::CapacityExceeded);
        }
        let dims = self.shape();
        let axis_len = dims[axis];
        let mut new_dims: Vec<usize> = dims
            .iter()
            .enumerate()
            .filter(|&(i, _)| i != axis)
            .map(|(_, &d)| d)
            .collect();
        if new_dims.is_empty() {
            new_dims.push(1);
        }
        let new_numel: usize = new_dims.iter().product();
        let mut result = vec![0.0f32; new_numel];

        let outer_size: usize = dims[..axis].iter().product();
        let inner_size: usize = dims[axis + 1..].iter().product();

        for outer in 0..outer_size {
            for inner in 0..inner_size {
                let mut sum = 0.0f32;
                for a in 0..axis_len {
                    let src_idx = outer * (axis_len * inner_size) + a * inner_size + inner;
                    sum += self.data[src_idx];
                }
                let dst_idx = outer * inner_size + inner;
                result[dst_idx] = sum;
            }
        }

        FabricTensor::from_slice(&new_dims, &result)
    }

    /// Elementwise sigmoid: `1 / (1 + exp(-x))`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[3], &[0.0, 100.0, -100.0]).unwrap();
    /// let b = a.sigmoid().unwrap();
    /// assert!((b.as_slice()[0] - 0.5).abs() < 1e-6);
    /// assert!((b.as_slice()[1] - 1.0).abs() < 1e-4);
    /// assert!(b.as_slice()[2] < 1e-4);
    /// ```
    pub fn sigmoid(&self) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_sigmoid")?;
            let cpu_result = self.sigmoid_cpu()?;
            return cpu_result.to_gpu();
        }
        self.sigmoid_cpu()
    }

    fn sigmoid_cpu(&self) -> Result<FabricTensor> {
        let data: Vec<f32> = self
            .data
            .iter()
            .map(|&x| 1.0 / (1.0 + (-x).exp()))
            .collect();
        FabricTensor::from_slice(self.shape(), &data)
    }

    /// Elementwise natural logarithm.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[3], &[1.0, core::f32::consts::E, 10.0]).unwrap();
    /// let b = a.ln().unwrap();
    /// assert!((b.as_slice()[0] - 0.0).abs() < 1e-6);
    /// assert!((b.as_slice()[1] - 1.0).abs() < 1e-6);
    /// ```
    pub fn ln(&self) -> Result<FabricTensor> {
        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_ln")?;
            let cpu_result = self.ln_cpu()?;
            return cpu_result.to_gpu();
        }
        self.ln_cpu()
    }

    fn ln_cpu(&self) -> Result<FabricTensor> {
        let data: Vec<f32> = self.data.iter().map(|&x| x.ln()).collect();
        FabricTensor::from_slice(self.shape(), &data)
    }

    /// Elementwise clamp to `[min, max]`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use grafos_tensor::FabricTensor;
    ///
    /// # grafos_std::host::reset_mock();
    /// # grafos_std::host::mock_set_fbmu_arena_size(65536);
    /// let a = FabricTensor::from_slice(&[4], &[-1.0, 0.5, 1.5, 3.0]).unwrap();
    /// let b = a.clip(0.0, 1.0).unwrap();
    /// assert_eq!(b.as_slice(), &[0.0, 0.5, 1.0, 1.0]);
    /// ```
    pub fn clip(&self, min: f32, max: f32) -> Result<FabricTensor> {
        let data: Vec<f32> = self.data.iter().map(|&x| x.clamp(min, max)).collect();
        FabricTensor::from_slice(self.shape(), &data)
    }

    /// Forward FFT of a 1-D real-valued tensor.
    ///
    /// Input must be a 1-D tensor with a power-of-2 number of elements.
    /// Returns a 1-D tensor of length `2*N` containing interleaved `[re, im]`
    /// pairs for each frequency bin.
    ///
    /// On GPU-placed tensors (with `gpu` feature), dispatches through the GPU
    /// submit path before computing the CPU result.
    pub fn fft(&self) -> Result<FabricTensor> {
        if self.ndim() != 1 {
            return Err(FabricError::CapacityExceeded);
        }
        let n = self.shape.dims[0];
        if n == 0 || !n.is_power_of_two() {
            return Err(FabricError::CapacityExceeded);
        }

        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_fft")?;
            let cpu_result = self.fft_cpu()?;
            return cpu_result.to_gpu();
        }

        self.fft_cpu()
    }

    /// Inverse FFT returning a 1-D real-valued tensor.
    ///
    /// Input must be a 1-D tensor of length `2*N` containing interleaved
    /// `[re, im]` pairs. `N` must be a power of 2. Returns a 1-D tensor
    /// of `N` real samples.
    ///
    /// On GPU-placed tensors (with `gpu` feature), dispatches through the GPU
    /// submit path before computing the CPU result.
    pub fn ifft(&self) -> Result<FabricTensor> {
        if self.ndim() != 1 {
            return Err(FabricError::CapacityExceeded);
        }
        let len = self.shape.dims[0];
        if len == 0 || !len.is_multiple_of(2) {
            return Err(FabricError::CapacityExceeded);
        }
        let n = len / 2;
        if !n.is_power_of_two() {
            return Err(FabricError::CapacityExceeded);
        }

        #[cfg(feature = "gpu")]
        if self.is_gpu() {
            self.dispatch_gpu_unary("tensor_ifft")?;
            let cpu_result = self.ifft_cpu()?;
            return cpu_result.to_gpu();
        }

        self.ifft_cpu()
    }

    fn fft_cpu(&self) -> Result<FabricTensor> {
        let n = self.shape.dims[0];
        // Build complex input from real data
        let mut complex = vec![(0.0f32, 0.0f32); n];
        for (i, &val) in self.data.iter().enumerate() {
            complex[i] = (val, 0.0);
        }

        // Bit-reverse permutation
        let bits = n.trailing_zeros();
        for i in 0..n {
            let j = i.reverse_bits() >> (usize::BITS - bits);
            if i < j {
                complex.swap(i, j);
            }
        }

        // Cooley-Tukey butterfly (forward)
        let mut len = 2;
        while len <= n {
            let half = len / 2;
            let angle = -2.0 * core::f32::consts::PI / len as f32;
            for start in (0..n).step_by(len) {
                let mut w_re = 1.0f32;
                let mut w_im = 0.0f32;
                let step_re = angle.cos();
                let step_im = angle.sin();
                for k in 0..half {
                    let (e_re, e_im) = complex[start + k];
                    let (o_re, o_im) = complex[start + k + half];
                    let tw_re = o_re * w_re - o_im * w_im;
                    let tw_im = o_re * w_im + o_im * w_re;
                    complex[start + k] = (e_re + tw_re, e_im + tw_im);
                    complex[start + k + half] = (e_re - tw_re, e_im - tw_im);
                    let new_w_re = w_re * step_re - w_im * step_im;
                    let new_w_im = w_re * step_im + w_im * step_re;
                    w_re = new_w_re;
                    w_im = new_w_im;
                }
            }
            len *= 2;
        }

        // Pack as interleaved [re, im, re, im, ...]
        let mut out = Vec::with_capacity(n * 2);
        for &(re, im) in &complex {
            out.push(re);
            out.push(im);
        }
        FabricTensor::from_slice(&[n * 2], &out)
    }

    fn ifft_cpu(&self) -> Result<FabricTensor> {
        let len = self.shape.dims[0];
        let n = len / 2;

        // Unpack interleaved [re, im, ...] into complex pairs
        let mut complex: Vec<(f32, f32)> = (0..n)
            .map(|i| (self.data[i * 2], self.data[i * 2 + 1]))
            .collect();

        // Bit-reverse permutation
        let bits = n.trailing_zeros();
        for i in 0..n {
            let j = i.reverse_bits() >> (usize::BITS - bits);
            if i < j {
                complex.swap(i, j);
            }
        }

        // Cooley-Tukey butterfly (inverse: positive angle)
        let mut blen = 2;
        while blen <= n {
            let half = blen / 2;
            let angle = 2.0 * core::f32::consts::PI / blen as f32;
            for start in (0..n).step_by(blen) {
                let mut w_re = 1.0f32;
                let mut w_im = 0.0f32;
                let step_re = angle.cos();
                let step_im = angle.sin();
                for k in 0..half {
                    let (e_re, e_im) = complex[start + k];
                    let (o_re, o_im) = complex[start + k + half];
                    let tw_re = o_re * w_re - o_im * w_im;
                    let tw_im = o_re * w_im + o_im * w_re;
                    complex[start + k] = (e_re + tw_re, e_im + tw_im);
                    complex[start + k + half] = (e_re - tw_re, e_im - tw_im);
                    let new_w_re = w_re * step_re - w_im * step_im;
                    let new_w_im = w_re * step_im + w_im * step_re;
                    w_re = new_w_re;
                    w_im = new_w_im;
                }
            }
            blen *= 2;
        }

        // Scale by 1/N and return real parts
        let scale = 1.0 / n as f32;
        let out: Vec<f32> = complex.iter().map(|&(re, _)| re * scale).collect();
        FabricTensor::from_slice(&[n], &out)
    }
}

/// Elementwise addition via the `+` operator.
///
/// Usage: `(&a + &b).unwrap()`
///
/// Delegates to [`FabricTensor::add`]. Both tensors must have the same shape.
impl<'a> Add for &'a FabricTensor {
    type Output = Result<FabricTensor>;

    fn add(self, rhs: &'a FabricTensor) -> Self::Output {
        self.add(rhs)
    }
}

/// Elementwise multiplication via the `*` operator.
///
/// Usage: `(&a * &b).unwrap()`
///
/// Delegates to [`FabricTensor::mul`]. Both tensors must have the same shape.
impl<'a> Mul for &'a FabricTensor {
    type Output = Result<FabricTensor>;

    fn mul(self, rhs: &'a FabricTensor) -> Self::Output {
        FabricTensor::mul(self, rhs)
    }
}

/// Scalar multiplication via the `*` operator.
///
/// Usage: `(&a * 2.0f32).unwrap()`
///
/// Delegates to [`FabricTensor::scale`].
impl Mul<f32> for &FabricTensor {
    type Output = Result<FabricTensor>;

    fn mul(self, rhs: f32) -> Self::Output {
        self.scale(rhs)
    }
}

/// Elementwise subtraction via the `-` operator.
///
/// Usage: `(&a - &b).unwrap()`
///
/// Delegates to [`FabricTensor::subtract`]. Both tensors must have the same shape.
impl<'a> Sub for &'a FabricTensor {
    type Output = Result<FabricTensor>;

    fn sub(self, rhs: &'a FabricTensor) -> Self::Output {
        self.subtract(rhs)
    }
}

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

    fn setup() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(1 << 20); // 1 MiB
    }

    // ---- Creation and element access ----

    #[test]
    fn zeros_creates_correct_shape() {
        setup();
        let t = FabricTensor::zeros(&[3, 4]).unwrap();
        assert_eq!(t.shape(), &[3, 4]);
        assert_eq!(t.ndim(), 2);
        assert_eq!(t.numel(), 12);
        assert_eq!(t.strides(), &[4, 1]);
        for i in 0..3 {
            for j in 0..4 {
                assert_eq!(t.get(&[i, j]).unwrap(), 0.0);
            }
        }
    }

    #[test]
    fn from_slice_roundtrip() {
        setup();
        let data: Vec<f32> = (1..=6).map(|x| x as f32).collect();
        let t = FabricTensor::from_slice(&[2, 3], &data).unwrap();
        assert_eq!(t.shape(), &[2, 3]);
        assert_eq!(t.get(&[0, 0]).unwrap(), 1.0);
        assert_eq!(t.get(&[0, 2]).unwrap(), 3.0);
        assert_eq!(t.get(&[1, 0]).unwrap(), 4.0);
        assert_eq!(t.get(&[1, 2]).unwrap(), 6.0);
    }

    #[test]
    fn from_slice_wrong_size() {
        setup();
        let result = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0]);
        assert!(result.is_err());
    }

    #[test]
    fn get_out_of_bounds() {
        setup();
        let t = FabricTensor::zeros(&[2, 3]).unwrap();
        assert!(t.get(&[2, 0]).is_err());
        assert!(t.get(&[0, 3]).is_err());
        assert!(t.get(&[0]).is_err()); // wrong ndim
    }

    #[test]
    fn set_element() {
        setup();
        let mut t = FabricTensor::zeros(&[2, 2]).unwrap();
        t.set(&[1, 0], 42.0).unwrap();
        assert_eq!(t.get(&[1, 0]).unwrap(), 42.0);
        assert_eq!(t.get(&[0, 0]).unwrap(), 0.0);
    }

    // ---- Matmul ----

    #[test]
    fn matmul_2x3_times_3x2() {
        setup();
        // A = [[1,2,3],[4,5,6]]
        let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        // B = [[7,8],[9,10],[11,12]]
        let b = FabricTensor::from_slice(&[3, 2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).unwrap();
        let c = a.matmul(&b).unwrap();
        assert_eq!(c.shape(), &[2, 2]);
        // C[0,0] = 1*7 + 2*9 + 3*11 = 58
        assert_eq!(c.get(&[0, 0]).unwrap(), 58.0);
        // C[0,1] = 1*8 + 2*10 + 3*12 = 64
        assert_eq!(c.get(&[0, 1]).unwrap(), 64.0);
        // C[1,0] = 4*7 + 5*9 + 6*11 = 139
        assert_eq!(c.get(&[1, 0]).unwrap(), 139.0);
        // C[1,1] = 4*8 + 5*10 + 6*12 = 154
        assert_eq!(c.get(&[1, 1]).unwrap(), 154.0);
    }

    #[test]
    fn matmul_incompatible_dims() {
        setup();
        let a = FabricTensor::zeros(&[2, 3]).unwrap();
        let b = FabricTensor::zeros(&[2, 3]).unwrap();
        assert!(a.matmul(&b).is_err());
    }

    #[test]
    fn matmul_not_2d() {
        setup();
        let a = FabricTensor::zeros(&[2, 3, 4]).unwrap();
        let b = FabricTensor::zeros(&[4, 2]).unwrap();
        assert!(a.matmul(&b).is_err());
    }

    // ---- Elementwise add ----

    #[test]
    fn add_elementwise() {
        setup();
        let a = FabricTensor::from_slice(&[2, 2], &[1.0, 2.0, 3.0, 4.0]).unwrap();
        let b = FabricTensor::from_slice(&[2, 2], &[10.0, 20.0, 30.0, 40.0]).unwrap();
        let c = a.add(&b).unwrap();
        assert_eq!(c.as_slice(), &[11.0, 22.0, 33.0, 44.0]);
    }

    #[test]
    fn add_shape_mismatch() {
        setup();
        let a = FabricTensor::zeros(&[2, 3]).unwrap();
        let b = FabricTensor::zeros(&[3, 2]).unwrap();
        assert!(a.add(&b).is_err());
    }

    // ---- Elementwise mul ----

    #[test]
    fn mul_elementwise() {
        setup();
        let a = FabricTensor::from_slice(&[3], &[2.0, 3.0, 4.0]).unwrap();
        let b = FabricTensor::from_slice(&[3], &[5.0, 6.0, 7.0]).unwrap();
        let c = FabricTensor::mul(&a, &b).unwrap();
        assert_eq!(c.as_slice(), &[10.0, 18.0, 28.0]);
    }

    // ---- Scale ----

    #[test]
    fn scale_scalar() {
        setup();
        let a = FabricTensor::from_slice(&[2, 2], &[1.0, 2.0, 3.0, 4.0]).unwrap();
        let b = a.scale(3.0).unwrap();
        assert_eq!(b.as_slice(), &[3.0, 6.0, 9.0, 12.0]);
    }

    // ---- ReLU ----

    #[test]
    fn relu_clamps_negatives() {
        setup();
        let a = FabricTensor::from_slice(&[4], &[-2.0, -0.5, 0.0, 3.0]).unwrap();
        let b = a.relu().unwrap();
        assert_eq!(b.as_slice(), &[0.0, 0.0, 0.0, 3.0]);
    }

    // ---- Softmax ----

    #[test]
    fn softmax_sums_to_one() {
        setup();
        let a = FabricTensor::from_slice(&[1, 4], &[1.0, 2.0, 3.0, 4.0]).unwrap();
        let b = a.softmax(1).unwrap();
        let sum: f32 = b.as_slice().iter().sum();
        assert!((sum - 1.0).abs() < 1e-6, "softmax sum = {sum}");
        // Values should be monotonically increasing.
        let s = b.as_slice();
        assert!(s[0] < s[1]);
        assert!(s[1] < s[2]);
        assert!(s[2] < s[3]);
    }

    #[test]
    fn softmax_2d_along_axis0() {
        setup();
        // 2x3 matrix, softmax along axis 0 (column-wise).
        let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let b = a.softmax(0).unwrap();
        let s = b.as_slice();
        // Each column should sum to 1.
        for col in 0..3 {
            let col_sum = s[col] + s[3 + col];
            assert!((col_sum - 1.0).abs() < 1e-6, "column {col} sum = {col_sum}");
        }
    }

    #[test]
    fn softmax_invalid_axis() {
        setup();
        let a = FabricTensor::zeros(&[2, 3]).unwrap();
        assert!(a.softmax(2).is_err());
    }

    // ---- Transpose ----

    #[test]
    fn transpose_2d() {
        setup();
        // [[1,2,3],[4,5,6]] -> [[1,4],[2,5],[3,6]]
        let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let b = a.transpose().unwrap();
        assert_eq!(b.shape(), &[3, 2]);
        assert_eq!(b.get(&[0, 0]).unwrap(), 1.0);
        assert_eq!(b.get(&[0, 1]).unwrap(), 4.0);
        assert_eq!(b.get(&[1, 0]).unwrap(), 2.0);
        assert_eq!(b.get(&[1, 1]).unwrap(), 5.0);
        assert_eq!(b.get(&[2, 0]).unwrap(), 3.0);
        assert_eq!(b.get(&[2, 1]).unwrap(), 6.0);
    }

    #[test]
    fn transpose_1d_fails() {
        setup();
        let a = FabricTensor::zeros(&[5]).unwrap();
        assert!(a.transpose().is_err());
    }

    // ---- Reshape ----

    #[test]
    fn reshape_preserves_data() {
        setup();
        let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let b = a.reshape(&[3, 2]).unwrap();
        assert_eq!(b.shape(), &[3, 2]);
        assert_eq!(b.as_slice(), a.as_slice());
    }

    #[test]
    fn reshape_wrong_numel() {
        setup();
        let a = FabricTensor::zeros(&[2, 3]).unwrap();
        assert!(a.reshape(&[2, 2]).is_err());
    }

    // ---- Operator overloading ----

    #[test]
    fn op_add() {
        setup();
        let a = FabricTensor::from_slice(&[3], &[1.0, 2.0, 3.0]).unwrap();
        let b = FabricTensor::from_slice(&[3], &[4.0, 5.0, 6.0]).unwrap();
        let c = (&a + &b).unwrap();
        assert_eq!(c.as_slice(), &[5.0, 7.0, 9.0]);
    }

    #[test]
    fn op_mul_elementwise() {
        setup();
        let a = FabricTensor::from_slice(&[3], &[2.0, 3.0, 4.0]).unwrap();
        let b = FabricTensor::from_slice(&[3], &[5.0, 6.0, 7.0]).unwrap();
        let c = (&a * &b).unwrap();
        assert_eq!(c.as_slice(), &[10.0, 18.0, 28.0]);
    }

    #[test]
    fn op_mul_scalar() {
        setup();
        let a = FabricTensor::from_slice(&[2], &[3.0, 4.0]).unwrap();
        let b = (&a * 2.0).unwrap();
        assert_eq!(b.as_slice(), &[6.0, 8.0]);
    }

    // ---- 3D tensor ----

    #[test]
    fn tensor_3d_access() {
        setup();
        // 2x3x4 tensor
        let data: Vec<f32> = (0..24).map(|x| x as f32).collect();
        let t = FabricTensor::from_slice(&[2, 3, 4], &data).unwrap();
        assert_eq!(t.ndim(), 3);
        assert_eq!(t.numel(), 24);
        assert_eq!(t.strides(), &[12, 4, 1]);
        // t[1][2][3] = 1*12 + 2*4 + 3 = 23
        assert_eq!(t.get(&[1, 2, 3]).unwrap(), 23.0);
    }

    // ---- Scalar (0-D) ----

    #[test]
    fn scalar_tensor() {
        setup();
        let t = FabricTensor::from_slice(&[], &[42.0]).unwrap();
        assert_eq!(t.ndim(), 0);
        assert_eq!(t.numel(), 1);
        assert_eq!(t.get(&[]).unwrap(), 42.0);
    }

    #[test]
    fn from_mem_lease_wraps_existing() {
        setup();
        let lease = MemBuilder::new().min_bytes(16).acquire().unwrap();
        let t = FabricTensor::from_mem_lease(&[2, 2], lease);
        assert_eq!(t.shape(), &[2, 2]);
        assert_eq!(t.numel(), 4);
        assert!(t.is_cpu());
    }

    // ---- 3D transpose ----

    #[test]
    fn transpose_3d_swaps_last_two() {
        setup();
        // Shape [2,3,4] -> [2,4,3]
        let data: Vec<f32> = (0..24).map(|x| x as f32).collect();
        let t = FabricTensor::from_slice(&[2, 3, 4], &data).unwrap();
        let t2 = t.transpose().unwrap();
        assert_eq!(t2.shape(), &[2, 4, 3]);
        // t[0][1][2] = 0*12 + 1*4 + 2 = 6 -> t2[0][2][1] = 6
        assert_eq!(t2.get(&[0, 2, 1]).unwrap(), 6.0);
        // t[1][0][3] = 1*12 + 0*4 + 3 = 15 -> t2[1][3][0] = 15
        assert_eq!(t2.get(&[1, 3, 0]).unwrap(), 15.0);
    }

    #[test]
    fn default_device_is_cpu() {
        setup();
        let t = FabricTensor::zeros(&[2, 2]).unwrap();
        assert_eq!(t.device(), Device::Cpu);
        assert!(t.is_cpu());
        assert!(!t.is_gpu());
    }

    #[test]
    #[cfg(not(feature = "gpu"))]
    fn to_gpu_without_feature_is_unsupported() {
        setup();
        let t = FabricTensor::zeros(&[2, 2]).unwrap();
        assert!(matches!(t.to_gpu(), Err(FabricError::Unsupported)));
    }

    #[test]
    #[cfg(feature = "gpu")]
    fn to_gpu_to_cpu_roundtrip_preserves_data() {
        setup();
        let t = FabricTensor::from_slice(&[2, 2], &[1.0, 2.0, 3.0, 4.0]).unwrap();
        let gpu_t = t.to_gpu().unwrap();
        assert!(gpu_t.is_gpu());
        assert!(!gpu_t.is_cpu());

        let cpu_t = gpu_t.to_cpu().unwrap();
        assert!(cpu_t.is_cpu());
        assert_eq!(cpu_t.shape(), &[2, 2]);
        assert_eq!(cpu_t.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    #[cfg(feature = "gpu")]
    fn gpu_matmul_dispatches_when_both_operands_are_gpu() {
        setup();
        let a = FabricTensor::from_slice(&[2, 2], &[1.0, 2.0, 3.0, 4.0])
            .unwrap()
            .to_gpu()
            .unwrap();
        let b = FabricTensor::from_slice(&[2, 2], &[5.0, 6.0, 7.0, 8.0])
            .unwrap()
            .to_gpu()
            .unwrap();
        host::test_mock::_set_gpu_session_error(Some(-1));
        assert!(matches!(a.matmul(&b), Err(FabricError::Disconnected)));
        host::test_mock::_set_gpu_session_error(None);
    }

    #[test]
    #[cfg(feature = "gpu")]
    fn gpu_binary_ops_return_gpu_placed_result() {
        setup();
        let a = FabricTensor::from_slice(&[2, 2], &[1.0, 2.0, 3.0, 4.0])
            .unwrap()
            .to_gpu()
            .unwrap();
        let b = FabricTensor::from_slice(&[2, 2], &[5.0, 6.0, 7.0, 8.0])
            .unwrap()
            .to_gpu()
            .unwrap();
        let result = a.add(&b).unwrap();
        assert!(result.is_gpu());
        assert_eq!(result.shape(), &[2, 2]);
        assert_eq!(result.as_slice(), &[6.0, 8.0, 10.0, 12.0]);
    }

    #[test]
    #[cfg(feature = "gpu")]
    fn gpu_unary_ops_return_gpu_placed_result() {
        setup();
        let a = FabricTensor::from_slice(&[4], &[-2.0, -0.5, 0.0, 3.0])
            .unwrap()
            .to_gpu()
            .unwrap();
        let result = a.relu().unwrap();
        assert!(result.is_gpu());
        assert_eq!(result.shape(), &[4]);
        assert_eq!(result.as_slice(), &[0.0, 0.0, 0.0, 3.0]);
    }

    // ---- FFT / IFFT ----

    #[test]
    fn fft_ifft_roundtrip() {
        setup();
        let input = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let t = FabricTensor::from_slice(&[8], &input).unwrap();
        let freq = t.fft().unwrap();
        assert_eq!(freq.shape(), &[16]); // 8 complex bins = 16 floats
        let recovered = freq.ifft().unwrap();
        assert_eq!(recovered.shape(), &[8]);
        for (i, &val) in recovered.as_slice().iter().enumerate() {
            assert!(
                (val - input[i]).abs() < 1e-4,
                "sample {i}: expected {}, got {val}",
                input[i]
            );
        }
    }

    #[test]
    fn fft_not_1d_fails() {
        setup();
        let t = FabricTensor::zeros(&[2, 4]).unwrap();
        assert!(t.fft().is_err());
    }

    #[test]
    fn fft_non_power_of_two_fails() {
        setup();
        let t = FabricTensor::from_slice(&[3], &[1.0, 2.0, 3.0]).unwrap();
        assert!(t.fft().is_err());
    }

    #[test]
    #[cfg(feature = "gpu")]
    fn gpu_fft_dispatches_and_produces_correct_result() {
        setup();
        let input = vec![1.0, 0.0, -1.0, 0.0];
        let t = FabricTensor::from_slice(&[4], &input)
            .unwrap()
            .to_gpu()
            .unwrap();
        let freq = t.fft().unwrap();
        assert!(freq.is_gpu());
        assert_eq!(freq.shape(), &[8]); // 4 bins * 2 floats

        // Verify roundtrip through IFFT
        let recovered = freq.ifft().unwrap();
        assert!(recovered.is_gpu());
        assert_eq!(recovered.shape(), &[4]);
        for (i, &val) in recovered.as_slice().iter().enumerate() {
            assert!(
                (val - input[i]).abs() < 1e-4,
                "sample {i}: expected {}, got {val}",
                input[i]
            );
        }
    }

    // ---- Subtract ----

    #[test]
    fn subtract_elementwise() {
        setup();
        let a = FabricTensor::from_slice(&[3], &[5.0, 3.0, 1.0]).unwrap();
        let b = FabricTensor::from_slice(&[3], &[1.0, 2.0, 3.0]).unwrap();
        let c = a.subtract(&b).unwrap();
        assert_eq!(c.as_slice(), &[4.0, 1.0, -2.0]);
    }

    #[test]
    fn subtract_operator() {
        setup();
        let a = FabricTensor::from_slice(&[2], &[10.0, 5.0]).unwrap();
        let b = FabricTensor::from_slice(&[2], &[3.0, 7.0]).unwrap();
        let c = (&a - &b).unwrap();
        assert_eq!(c.as_slice(), &[7.0, -2.0]);
    }

    #[test]
    fn subtract_shape_mismatch() {
        setup();
        let a = FabricTensor::from_slice(&[2], &[1.0, 2.0]).unwrap();
        let b = FabricTensor::from_slice(&[3], &[1.0, 2.0, 3.0]).unwrap();
        assert!(a.subtract(&b).is_err());
    }

    // ---- Sum axis ----

    #[test]
    fn sum_axis_row_sums() {
        setup();
        let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let s = a.sum_axis(1).unwrap();
        assert_eq!(s.shape(), &[2]);
        assert_eq!(s.as_slice(), &[6.0, 15.0]);
    }

    #[test]
    fn sum_axis_col_sums() {
        setup();
        let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let s = a.sum_axis(0).unwrap();
        assert_eq!(s.shape(), &[3]);
        assert_eq!(s.as_slice(), &[5.0, 7.0, 9.0]);
    }

    #[test]
    fn sum_axis_1d() {
        setup();
        let a = FabricTensor::from_slice(&[4], &[1.0, 2.0, 3.0, 4.0]).unwrap();
        let s = a.sum_axis(0).unwrap();
        assert_eq!(s.shape(), &[1]);
        assert_eq!(s.as_slice(), &[10.0]);
    }

    #[test]
    fn sum_axis_out_of_bounds() {
        setup();
        let a = FabricTensor::from_slice(&[2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        assert!(a.sum_axis(2).is_err());
    }

    // ---- Sigmoid ----

    #[test]
    fn sigmoid_values() {
        setup();
        let a = FabricTensor::from_slice(&[3], &[0.0, 100.0, -100.0]).unwrap();
        let b = a.sigmoid().unwrap();
        assert!((b.as_slice()[0] - 0.5).abs() < 1e-6);
        assert!((b.as_slice()[1] - 1.0).abs() < 1e-4);
        assert!(b.as_slice()[2] < 1e-4);
    }

    // ---- Ln ----

    #[test]
    fn ln_values() {
        setup();
        let a = FabricTensor::from_slice(&[3], &[1.0, core::f32::consts::E, 10.0]).unwrap();
        let b = a.ln().unwrap();
        assert!((b.as_slice()[0] - 0.0).abs() < 1e-6);
        assert!((b.as_slice()[1] - 1.0).abs() < 1e-5);
        assert!((b.as_slice()[2] - 10.0f32.ln()).abs() < 1e-5);
    }

    // ---- Clip ----

    #[test]
    fn clip_values() {
        setup();
        let a = FabricTensor::from_slice(&[4], &[-1.0, 0.5, 1.5, 3.0]).unwrap();
        let b = a.clip(0.0, 1.0).unwrap();
        assert_eq!(b.as_slice(), &[0.0, 0.5, 1.0, 1.0]);
    }
}