torsh_tensor/

math_ops.rs

//! Mathematical operations for tensors - Enhanced with SciRS2 Performance Features
//!
//! This module provides comprehensive mathematical operations including arithmetic,
//! trigonometric, exponential, and logarithmic functions, along with scalar operations
//! and **full SciRS2 backend integration** for optimal performance.
//!
//! # Enhanced SciRS2 Features
//!
//! - **SIMD Acceleration**: Vectorized operations using SciRS2 SIMD support
//! - **Parallel Processing**: Multi-core tensor operations via SciRS2 parallel framework
//! - **GPU Acceleration**: CUDA/Metal backend support through SciRS2 GPU abstraction
//! - **Memory Optimization**: Lazy evaluation and memory-mapped arrays for large tensors
//! - **Scalar arithmetic**: Operations between tensors and scalar values
//! - **Element-wise operations**: Basic arithmetic operations between tensors
//! - **Mathematical functions**: sqrt, exp, log, trigonometric functions
//! - **Complex operations**: Support for complex number operations
//! - **Broadcasting**: Automatic shape compatibility for operations

use std::sync::Arc;
use torsh_core::{
    dtype::TensorElement,
    error::{Result, TorshError},
};

use crate::memory_pool::global_acquire_uninit;

// ✅ SciRS2 Advanced Features Integration
// Performance acceleration through SciRS2 ecosystem
#[cfg(feature = "simd")]
mod simd_imports {
    // ✅ SciRS2 Breakthrough SIMD Implementation - 14.17x Performance

    // 🚀 Hyperoptimized SIMD implementations with breakthrough performance
    // Note: Currently not using direct SIMD functions - they're available via SimdUnifiedOps trait
    // pub use scirs2_core::simd::{ ... };

    // Array types
    pub use scirs2_core::ndarray::Array1;
}

#[cfg(feature = "simd")]
use simd_imports::*;

// Chunking and parallel processing
#[cfg(feature = "parallel")]
use scirs2_core::chunking::{
    CacheAwareness, ChunkConfig, ChunkStrategy, ComputeIntensity, GpuChunkSettings, MemoryPattern,
    NumaStrategy,
};

// Memory optimization features
// Note: memory_efficient features require enabling the memory_efficient feature flag
// use scirs2_core::memory_efficient::{MemoryMappedArray, LazyArray};

// TODO: scirs2_core::profiling module not available yet
// Performance profiling integration
// #[cfg(feature = "profiling")]
// use scirs2_core::profiling::Profiler;
// TODO: profile_section macro not available in scirs2_core yet
// use scirs2_core::profiling::profile_section;

use crate::core_ops::{Operation, Tensor, UnaryKind};
#[cfg(feature = "simd")]
use crate::simd_ops_f32::BinaryF32Op;
use crate::storage::TensorStorage;

impl<T: TensorElement> Tensor<T> {
    /// Record `result` as a differentiable element-wise unary of `self`.
    ///
    /// Called at the end of the public activation/transcendental wrappers so
    /// every dispatch path (SIMD, parallel, GPU, scalar) shares one recorded
    /// derivative. A no-op when gradients are not being tracked.
    pub(crate) fn record_unary(&self, mut result: Self, kind: UnaryKind) -> Self {
        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::Unary {
                input: Arc::new(self.clone()),
                kind,
            };
        }
        result
    }

    /// Record `result` as `leaky_relu(self, negative_slope)`.
    ///
    /// Parameterized twin of [`Tensor::record_unary`]: the slope cannot live
    /// in [`UnaryKind`], so it goes in its own `Operation` variant (the
    /// `MulScalar`/`DivScalar` idiom). A no-op when gradients are not tracked.
    pub(crate) fn record_leaky_relu(&self, mut result: Self, negative_slope: T) -> Self {
        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::LeakyRelu {
                input: Arc::new(self.clone()),
                negative_slope,
            };
        }
        result
    }

    /// Record `result` as `maximum(self, other)` (`is_maximum`) or
    /// `minimum(self, other)`.
    ///
    /// The operands are stored *un-broadcast*, like [`Operation::Add`]'s, so
    /// backward folds each gradient back to its own operand's shape.
    pub(crate) fn record_extremum(&self, other: &Self, mut result: Self, is_maximum: bool) -> Self {
        if crate::should_record_grad(self.requires_grad || other.requires_grad) {
            result.requires_grad = true;
            let lhs = Arc::new(self.clone());
            let rhs = Arc::new(other.clone());
            result.operation = if is_maximum {
                Operation::Maximum { lhs, rhs }
            } else {
                Operation::Minimum { lhs, rhs }
            };
        }
        result
    }

    /// Record `result` as `abs(self)` — but only when `self` and `result` have
    /// the *same* element type.
    ///
    /// `Tensor::abs` lives on the [`torsh_core::dtype::ComplexElement`] impl
    /// block (`complex_ops.rs`) and has signature
    /// `Tensor<T> -> Tensor<T::Real>`. Real element types implement
    /// `ComplexElement` with `Real = Self`, so for them the two types coincide
    /// and `|x|` is the ordinary real absolute value with sub-gradient
    /// `sign(x)`. For a genuinely complex `T` the output type differs, the
    /// derivative is the Wirtinger `z / |z|` rather than `sign(x)`, and
    /// [`Operation::Unary`] could not hold the operand anyway (it wants an
    /// `Arc<Tensor<R>>`, not an `Arc<Tensor<T>>`).
    ///
    /// The `Any` downcast is what separates the two cases: it succeeds exactly
    /// when `T == R`, so complex `abs` keeps its historical detached behaviour
    /// and real `abs` records. It is a safe downcast — `TensorElement: 'static`
    /// supplies the `Any` bound — not a transmute.
    pub(crate) fn record_abs_if_real<R: TensorElement>(&self, result: Tensor<R>) -> Tensor<R> {
        match (self as &dyn std::any::Any).downcast_ref::<Tensor<R>>() {
            Some(real_self) => real_self.record_unary(result, UnaryKind::Abs),
            None => result,
        }
    }
}

/// Element count from which a same-shape binary op dispatches to the hardware
/// SIMD kernels of `scirs2_core::simd_ops::SimdUnifiedOps`.
///
/// Below this size the per-call dispatch and buffer acquisition dominate the
/// arithmetic, so the generic path is faster.
#[cfg(feature = "simd")]
const SIMD_BINARY_THRESHOLD: usize = 1024;

/// Element count from which the generic element-wise path fans out over
/// `scirs2_core::parallel_ops` worker threads.
///
/// The same value as [`SIMD_BINARY_THRESHOLD`] on purpose: it leaves no size
/// band in which neither the vector kernels nor the worker pool apply.
#[cfg(feature = "parallel")]
const PARALLEL_ELEMENTWISE_THRESHOLD: usize = 1024;

/// Dispatch one of the f64 vector kernels of
/// `scirs2_core::simd_ops::SimdUnifiedOps` into a caller-supplied buffer.
///
/// The op tag is shared with the f32 kernels (`BinaryF32Op`); only the element
/// type differs. All three slices must have the same length.
#[cfg(feature = "simd")]
fn dispatch_f64_into(op: BinaryF32Op, a: &[f64], b: &[f64], out: &mut [f64]) {
    use scirs2_core::simd_ops::SimdUnifiedOps;
    match op {
        BinaryF32Op::Add => <f64 as SimdUnifiedOps>::simd_add_into(a, b, out),
        BinaryF32Op::Sub => <f64 as SimdUnifiedOps>::simd_sub_into(a, b, out),
        BinaryF32Op::Mul => <f64 as SimdUnifiedOps>::simd_mul_into(a, b, out),
        BinaryF32Op::Div => <f64 as SimdUnifiedOps>::simd_div_into(a, b, out),
    }
}

// 🚀 Adaptive SIMD Selection System for Maximum Performance
#[cfg(feature = "simd")]
pub(crate) mod adaptive_simd {
    use super::*;
    use scirs2_core::ndarray::ArrayView1;

    /// SIMD-accelerated ReLU activation function
    /// Uses scirs2_core SIMD implementation for optimal performance
    pub fn adaptive_simd_relu_f32(input: &ArrayView1<f32>) -> Array1<f32> {
        scirs2_core::simd::activation::simd_relu_f32(input)
    }

    /// SIMD-accelerated sigmoid activation function
    /// Uses scirs2_core SIMD implementation for optimal performance
    pub fn adaptive_simd_sigmoid_f32(input: &ArrayView1<f32>) -> Array1<f32> {
        scirs2_core::simd::transcendental::simd_sigmoid_f32(input)
    }

    // NOTE: there is deliberately no `adaptive_simd_gelu_f32` shim.
    // `scirs2_core::simd::transcendental::simd_gelu_f32` computes a
    // clamped-Pade approximation of `tanh`, i.e. a materially different
    // function from the exact-tanh GELU the crate documents and
    // differentiates (measured: up to 0.0211 absolute, 515% relative at
    // x = -2.40). Routing large f32 tensors through it made `Tensor::gelu`
    // discontinuous across its own `numel = 1000` dispatch threshold and
    // inconsistent with its recorded backward, so the shim and its dispatch
    // were removed rather than documented around. See
    // `Tensor::gelu_forward` in math_ops_trig.rs. `relu` and `sigmoid` are
    // unaffected: their SIMD kernels match their formulas exactly.
}

#[cfg(feature = "simd")]
// 🚀 Intelligent Chunking System for Advanced Optimization
#[cfg(feature = "parallel")]
mod intelligent_chunking {
    use super::*;

    /// Tensor operation types for optimal chunking strategy selection
    #[derive(Debug, Clone, Copy)]
    pub enum TensorOpType {
        /// Element-wise operations (add, mul, etc.)
        ElementWise,
        /// Activation functions
        Activation,
    }

    // Note: GpuChunkSettings is now imported from scirs2_core::chunking

    /// Create optimal chunking configuration based on tensor operation characteristics
    pub fn create_optimal_chunk_config(
        tensor_size: usize,
        op_type: TensorOpType,
        _device: torsh_core::device::DeviceType,
        is_gpu_available: bool,
    ) -> ChunkConfig {
        match op_type {
            TensorOpType::ElementWise => ChunkConfig {
                strategy: if tensor_size > 100_000 {
                    ChunkStrategy::MemoryOptimized
                } else {
                    ChunkStrategy::CacheOptimized
                },
                min_chunk_size: 64,
                max_chunk_size: 8192,
                prefer_work_stealing: true,
                memory_pattern: MemoryPattern::Sequential,
                compute_intensity: ComputeIntensity::MemoryBound,
                enable_monitoring: false,
                load_balance_factor: 0.1,
                cache_awareness: CacheAwareness::L2,
                numa_strategy: NumaStrategy::LocalPreferred,
                gpu_settings: if is_gpu_available {
                    Some(GpuChunkSettings::default())
                } else {
                    None
                },
            },

            TensorOpType::Activation => ChunkConfig {
                strategy: ChunkStrategy::CacheOptimized,
                min_chunk_size: 64,
                max_chunk_size: 4096,
                prefer_work_stealing: true,
                memory_pattern: MemoryPattern::Sequential,
                compute_intensity: ComputeIntensity::ComputeIntensive,
                enable_monitoring: false,
                load_balance_factor: 0.1,
                cache_awareness: CacheAwareness::L1,
                numa_strategy: NumaStrategy::LocalPreferred,
                gpu_settings: if is_gpu_available {
                    Some(GpuChunkSettings {
                        gpu_memory_ratio: 0.7,
                        gpu_min_chunk: 2048,
                        overlap_compute: true,
                        gpu_bandwidth: None,      // Option<u64>
                        transfer_bandwidth: None, // Option<u64>
                    })
                } else {
                    None
                },
            },
        }
    }

    /// Intelligent chunking for parallel tensor operations
    pub fn intelligent_parallel_process<T, F, R>(
        data: Vec<T>,
        op_type: TensorOpType,
        device: torsh_core::device::DeviceType,
        operation: F,
    ) -> Vec<R>
    where
        T: Send + Sync,
        R: Send + Sync,
        F: Fn(T) -> R + Send + Sync,
    {
        let is_gpu_available = matches!(
            device,
            torsh_core::device::DeviceType::Cuda(_)
                | torsh_core::device::DeviceType::Metal(_)
                | torsh_core::device::DeviceType::Wgpu(_)
        );

        let _chunk_config =
            create_optimal_chunk_config(data.len(), op_type, device, is_gpu_available);

        // ✅ SciRS2 POLICY: Use scirs2_core::parallel_ops instead of direct rayon
        #[cfg(feature = "parallel")]
        {
            use scirs2_core::parallel_ops::*;
            data.into_par_iter().map(operation).collect()
        }
        #[cfg(not(feature = "parallel"))]
        {
            data.into_iter().map(operation).collect()
        }
    }
}

#[cfg(feature = "parallel")]
use intelligent_chunking::*;

/// Check if two shapes can be broadcasted together
fn can_broadcast(shape1: &[usize], shape2: &[usize]) -> bool {
    let max_dims = shape1.len().max(shape2.len());

    for i in 0..max_dims {
        let dim1 = if i < shape1.len() {
            shape1[shape1.len() - 1 - i]
        } else {
            1
        };
        let dim2 = if i < shape2.len() {
            shape2[shape2.len() - 1 - i]
        } else {
            1
        };

        if dim1 != dim2 && dim1 != 1 && dim2 != 1 {
            return false;
        }
    }
    true
}

/// Compute the resulting shape after broadcasting
fn compute_broadcast_shape(shape1: &[usize], shape2: &[usize]) -> Result<Vec<usize>> {
    let max_dims = shape1.len().max(shape2.len());
    let mut result = Vec::with_capacity(max_dims);

    for i in 0..max_dims {
        let dim1 = if i < shape1.len() {
            shape1[shape1.len() - 1 - i]
        } else {
            1
        };
        let dim2 = if i < shape2.len() {
            shape2[shape2.len() - 1 - i]
        } else {
            1
        };

        if dim1 == dim2 {
            result.push(dim1);
        } else if dim1 == 1 {
            result.push(dim2);
        } else if dim2 == 1 {
            result.push(dim1);
        } else {
            return Err(TorshError::ShapeMismatch {
                expected: shape1.to_vec(),
                got: shape2.to_vec(),
            });
        }
    }

    result.reverse();
    Ok(result)
}

/// Compute the index in the original tensor given a flat index in the broadcasted tensor
fn compute_broadcast_index(
    flat_idx: usize,
    broadcast_shape: &[usize],
    original_shape: &[usize],
) -> usize {
    let mut result = 0;
    let mut remaining = flat_idx;

    let dims_diff = broadcast_shape.len() - original_shape.len();

    for (i, &broadcast_dim) in broadcast_shape.iter().enumerate() {
        // Calculate stride (product of remaining dimensions)
        let stride = broadcast_shape[i + 1..].iter().product::<usize>().max(1);
        let coord = remaining / stride;
        remaining %= stride;

        // Validate coordinate is within broadcast dimension
        debug_assert!(
            coord < broadcast_dim,
            "Coordinate {} out of bounds for dimension {} of size {}",
            coord,
            i,
            broadcast_dim
        );

        if i >= dims_diff {
            let original_dim = original_shape[i - dims_diff];
            let adjusted_coord = if original_dim == 1 { 0 } else { coord };
            result = result * original_dim + adjusted_coord;
        }
    }

    result
}

impl<T: TensorElement + Copy> Tensor<T> {
    /// Whether two tensors read through the very same storage allocation.
    ///
    /// Binary ops borrow both operands out of storage instead of copying them,
    /// and `RwLock::read` is not re-entrant: taking a guard on one buffer while
    /// already holding one on the *same* buffer risks a dead-lock. Callers use
    /// this to fall back to a materialised copy of the second operand — which
    /// is what makes `t.add(&t)` and `t.add(&t.clone())` safe.
    ///
    /// Unlike [`Tensor::shares_storage`] this covers every storage variant,
    /// including the two SIMD ones.
    pub(crate) fn shares_storage_buffer(&self, other: &Self) -> bool {
        match (&self.storage, &other.storage) {
            (TensorStorage::InMemory(a), TensorStorage::InMemory(b)) => Arc::ptr_eq(a, b),
            (TensorStorage::MemoryMapped(a), TensorStorage::MemoryMapped(b)) => Arc::ptr_eq(a, b),
            #[cfg(feature = "simd")]
            (TensorStorage::Aligned(a), TensorStorage::Aligned(b)) => Arc::ptr_eq(a, b),
            #[cfg(feature = "simd")]
            (TensorStorage::SimdOptimized(a), TensorStorage::SimdOptimized(b)) => Arc::ptr_eq(a, b),
            #[cfg(feature = "gpu")]
            (TensorStorage::Device { buffer: a, .. }, TensorStorage::Device { buffer: b, .. }) => {
                Arc::ptr_eq(a, b)
            }
            _ => false,
        }
    }

    /// Address of the storage allocation this tensor reads through.
    ///
    /// Used only to impose a total order on lock acquisition; the value is
    /// never dereferenced.
    fn storage_addr(&self) -> usize {
        match &self.storage {
            TensorStorage::InMemory(data) => Arc::as_ptr(data) as usize,
            TensorStorage::MemoryMapped(storage) => Arc::as_ptr(storage) as usize,
            #[cfg(feature = "simd")]
            TensorStorage::Aligned(data) => Arc::as_ptr(data) as usize,
            #[cfg(feature = "simd")]
            TensorStorage::SimdOptimized(storage) => Arc::as_ptr(storage) as usize,
            #[cfg(feature = "gpu")]
            TensorStorage::Device { buffer, .. } => Arc::as_ptr(buffer) as usize,
        }
    }

    /// Run `f` against contiguous, view-ordered slices of both operands.
    ///
    /// Neither operand is copied while they live in different buffers: each
    /// slice is borrowed straight out of storage for the duration of the
    /// closure. Strided views are materialised in view order first (so the
    /// closure can always index row-major).
    ///
    /// # Locking
    /// Two operands that share one allocation take a single copy of `other`
    /// rather than nesting guards on the same (non-re-entrant) lock, which is
    /// what makes `t.add(&t)` and `t.add(&t.clone())` safe. Distinct buffers
    /// are locked in ascending address order, so two threads running mirrored
    /// operations (`a op b` and `b op a`) can never build an acquisition cycle.
    pub(crate) fn with_operand_slices<R, F>(&self, other: &Self, f: F) -> Result<R>
    where
        F: FnOnce(&[T], &[T]) -> Result<R>,
    {
        if self.shares_storage_buffer(other) {
            // Materialise before acquiring this tensor's guard, never inside it.
            let other_data = other.to_vec()?;
            return self.with_contiguous_data(|lhs| f(lhs, &other_data));
        }
        if self.storage_addr() <= other.storage_addr() {
            self.with_contiguous_data(|lhs| other.with_contiguous_data(|rhs| f(lhs, rhs)))
        } else {
            other.with_contiguous_data(|rhs| self.with_contiguous_data(|lhs| f(lhs, rhs)))
        }
    }

    /// Try the hardware-SIMD fast path for a same-shape binary operation.
    ///
    /// Returns `Ok(None)` when the operands are not eligible — different
    /// shapes, fewer than [`SIMD_BINARY_THRESHOLD`] elements, or an element
    /// type without a vector kernel — in which case the caller falls back to
    /// [`Tensor::elementwise_operation`]. `f32` and `f64` both have real
    /// AVX2/NEON kernels in `scirs2_core::simd_ops`.
    ///
    /// Allocation profile: zero copies of the operands and exactly one pooled
    /// output buffer, which the kernel initialises in full.
    #[cfg(feature = "simd")]
    fn try_simd_binary(&self, other: &Self, op: BinaryF32Op) -> Result<Option<Self>> {
        let n = self.numel();
        if n < SIMD_BINARY_THRESHOLD || self.shape() != other.shape() {
            return Ok(None);
        }
        let is_f32 = std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>();
        let is_f64 = std::any::TypeId::of::<T>() == std::any::TypeId::of::<f64>();
        if !is_f32 && !is_f64 {
            return Ok(None);
        }

        let result_data = self.with_operand_slices(other, |a, b| {
            if a.len() != n || b.len() != n {
                return Err(TorshError::ShapeMismatch {
                    expected: vec![n],
                    got: vec![b.len()],
                });
            }
            let mut buf = global_acquire_uninit::<T>(n);
            {
                let uninit = &mut buf.as_uninit_slice_mut()[..n];
                // SAFETY: `MaybeUninit<T>` has the layout of `T`, and every
                // kernel invoked below is store-only — `simd_*_into` writes all
                // `n` outputs before anything reads them — so the buffer is
                // fully initialised when this block ends. That is exactly why
                // no zero-fill pass is needed here.
                let out: &mut [T] =
                    unsafe { std::slice::from_raw_parts_mut(uninit.as_mut_ptr().cast::<T>(), n) };
                if is_f32 {
                    // SAFETY: the TypeId check above confirms `T == f32`, so
                    // reinterpreting the same-layout slices is a no-op.
                    unsafe {
                        op.dispatch_into(
                            std::slice::from_raw_parts(a.as_ptr().cast::<f32>(), n),
                            std::slice::from_raw_parts(b.as_ptr().cast::<f32>(), n),
                            std::slice::from_raw_parts_mut(out.as_mut_ptr().cast::<f32>(), n),
                        );
                    }
                } else {
                    // SAFETY: the TypeId check above confirms `T == f64`.
                    unsafe {
                        dispatch_f64_into(
                            op,
                            std::slice::from_raw_parts(a.as_ptr().cast::<f64>(), n),
                            std::slice::from_raw_parts(b.as_ptr().cast::<f64>(), n),
                            std::slice::from_raw_parts_mut(out.as_mut_ptr().cast::<f64>(), n),
                        );
                    }
                }
            }
            Ok(buf.into_vec(n))
        })?;

        let result = Self::from_data(result_data, self.shape().dims().to_vec(), self.device)?;
        Ok(Some(result))
    }

    /// Add scalar to all elements in-place
    pub fn add_scalar_(&mut self, scalar: T) -> Result<()>
    where
        T: Copy + std::ops::Add<Output = T>,
    {
        // Ensure data is unique (copy-on-write)
        self.make_unique()?;
        self.apply_(|x| x + scalar)
    }

    /// Add scalar to all elements (returns new tensor)
    ///
    /// Records [`Operation::AddScalar`]: `map` is forward-only, so without this
    /// the shifted tensor would be a `requires_grad` leaf that swallows the
    /// gradient of everything upstream of it.
    pub fn add_scalar(&self, scalar: T) -> Result<Self>
    where
        T: Copy + std::ops::Add<Output = T>,
    {
        let mut result = self.map(|x| x + scalar)?;
        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::AddScalar {
                input: Arc::new(self.clone()),
                scalar,
            };
        }
        Ok(result)
    }

    /// Subtract scalar from all elements in-place
    pub fn sub_scalar_(&mut self, scalar: T) -> Result<()>
    where
        T: Copy + std::ops::Sub<Output = T>,
    {
        self.make_unique()?;
        self.apply_(|x| x - scalar)
    }

    /// Subtract scalar from all elements (returns new tensor)
    ///
    /// Records [`Operation::AddScalar`] with the negated constant — the
    /// backward rule is the same identity — so the public bound stays
    /// `Sub<Output = T>` and no `Neg`/`Add` bound leaks into the API.
    pub fn sub_scalar(&self, scalar: T) -> Result<Self>
    where
        T: Copy + std::ops::Sub<Output = T>,
    {
        let mut result = self.map(|x| x - scalar)?;
        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::AddScalar {
                input: Arc::new(self.clone()),
                scalar: <T as TensorElement>::zero() - scalar,
            };
        }
        Ok(result)
    }

    /// Multiply all elements by scalar in-place
    pub fn mul_scalar_(&mut self, scalar: T) -> Result<()>
    where
        T: Copy + std::ops::Mul<Output = T>,
    {
        // Ensure data is unique (copy-on-write)
        self.make_unique()?;
        self.apply_(|x| x * scalar)
    }

    /// Multiply all elements by scalar (returns new tensor)
    ///
    /// The result joins the autograd graph when the input requires gradients:
    /// `d/dinput (input * s) = s`. Recording a dedicated scalar node keeps the
    /// tape free of a full tensor of copies of the constant.
    pub fn mul_scalar(&self, scalar: T) -> Result<Self>
    where
        T: Copy + std::ops::Mul<Output = T>,
    {
        let mut result = self.map(|x| x * scalar)?;
        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::MulScalar {
                input: Arc::new(self.clone()),
                scalar,
            };
        }
        Ok(result)
    }

    /// Divide all elements by scalar in-place
    pub fn div_scalar_(&mut self, scalar: T) -> Result<()>
    where
        T: Copy + std::ops::Div<Output = T>,
    {
        self.make_unique()?;
        self.apply_(|x| x / scalar)
    }

    /// Divide all elements by scalar (returns new tensor)
    ///
    /// The result joins the autograd graph when the input requires gradients:
    /// `d/dinput (input / s) = 1 / s`.
    pub fn div_scalar(&self, scalar: T) -> Result<Self>
    where
        T: Copy + std::ops::Div<Output = T>,
    {
        let mut result = self.map(|x| x / scalar)?;
        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::DivScalar {
                input: Arc::new(self.clone()),
                scalar,
            };
        }
        Ok(result)
    }

    /// Element-wise addition with another tensor (supports broadcasting)
    pub fn add(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Add<Output = T>,
    {
        // Handle broadcasting for different shapes
        if self.shape() != other.shape() {
            return self.broadcast_add(other);
        }

        // Hardware SIMD fast path (f32/f64, ≥ SIMD_BINARY_THRESHOLD elements),
        // reading both operands in place and writing one pooled output buffer.
        #[cfg(feature = "simd")]
        if let Some(mut result) = self.try_simd_binary(other, BinaryF32Op::Add)? {
            if crate::should_record_grad(self.requires_grad || other.requires_grad) {
                result.requires_grad = true;
                result.operation = Operation::Add {
                    lhs: Arc::new(self.clone()),
                    rhs: Arc::new(other.clone()),
                };
            }
            return Ok(result);
        }

        // Same shape - use optimized elementwise operation
        let mut result = self.elementwise_operation(other, |a, b| a + b)?;

        // Preserve gradient tracking
        if crate::should_record_grad(self.requires_grad || other.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::Add {
                lhs: Arc::new(self.clone()),
                rhs: Arc::new(other.clone()),
            };
        }

        Ok(result)
    }

    /// Broadcasting-aware addition for different shapes
    fn broadcast_add(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Add<Output = T>,
    {
        // Simple broadcasting implementation for common cases
        let self_shape_binding = self.shape();
        let other_shape_binding = other.shape();
        let self_shape = self_shape_binding.dims();
        let other_shape = other_shape_binding.dims();

        // Check if broadcasting is possible
        if !can_broadcast(self_shape, other_shape) {
            return Err(TorshError::ShapeMismatch {
                expected: self_shape.to_vec(),
                got: other_shape.to_vec(),
            });
        }

        // Compute the broadcasted shape
        let broadcast_shape = compute_broadcast_shape(self_shape, other_shape)?;

        // Perform broadcasting addition, reading both operands in place.
        let total_elems: usize = broadcast_shape.iter().product();
        let result_data = self.with_operand_slices(other, |self_data, other_data| {
            let mut buf = global_acquire_uninit::<T>(total_elems);
            let uninit = &mut buf.as_uninit_slice_mut()[..total_elems];
            let mut count = 0;

            for i in 0..total_elems {
                let self_idx = compute_broadcast_index(i, &broadcast_shape, self_shape);
                let other_idx = compute_broadcast_index(i, &broadcast_shape, other_shape);

                let self_val = *self_data
                    .get(self_idx)
                    .ok_or_else(|| TorshError::IndexError {
                        index: self_idx,
                        size: self_data.len(),
                    })?;
                let other_val =
                    *other_data
                        .get(other_idx)
                        .ok_or_else(|| TorshError::IndexError {
                            index: other_idx,
                            size: other_data.len(),
                        })?;
                uninit[count].write(self_val + other_val);
                count += 1;
            }

            Ok(buf.into_vec(count))
        })?;
        let mut result = Self::from_data(result_data, broadcast_shape, self.device)?;

        // Preserve gradient tracking
        if crate::should_record_grad(self.requires_grad || other.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::Add {
                lhs: Arc::new(self.clone()),
                rhs: Arc::new(other.clone()),
            };
        }

        Ok(result)
    }

    /// Element-wise subtraction with another tensor
    pub fn sub(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Sub<Output = T>,
    {
        // Hardware SIMD fast path (f32/f64, ≥ SIMD_BINARY_THRESHOLD elements),
        // reading both operands in place and writing one pooled output buffer.
        #[cfg(feature = "simd")]
        if let Some(mut result) = self.try_simd_binary(other, BinaryF32Op::Sub)? {
            if crate::should_record_grad(self.requires_grad || other.requires_grad) {
                result.requires_grad = true;
                result.operation = crate::Operation::Sub {
                    lhs: Arc::new(self.clone()),
                    rhs: Arc::new(other.clone()),
                };
            }
            return Ok(result);
        }

        let mut result = self.elementwise_operation(other, |a, b| a - b)?;

        // Propagate requires_grad and record operation for autograd
        if crate::should_record_grad(self.requires_grad || other.requires_grad) {
            result.requires_grad = true;
            result.operation = crate::Operation::Sub {
                lhs: Arc::new(self.clone()),
                rhs: Arc::new(other.clone()),
            };
        }

        Ok(result)
    }

    /// Element-wise multiplication with another tensor (supports broadcasting).
    ///
    /// When either operand requires gradients the result records
    /// `Operation::Mul` with the *un-broadcast* operands, so backward can apply
    /// `d/dlhs = rhs`, `d/drhs = lhs` and fold each gradient back to its
    /// operand's own shape.
    pub fn mul(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Mul<Output = T>,
    {
        // Hardware SIMD fast path (f32/f64, ≥ SIMD_BINARY_THRESHOLD elements),
        // reading both operands in place and writing one pooled output buffer.
        #[cfg(feature = "simd")]
        if let Some(mut result) = self.try_simd_binary(other, BinaryF32Op::Mul)? {
            if crate::should_record_grad(self.requires_grad || other.requires_grad) {
                result.requires_grad = true;
                result.operation = crate::Operation::Mul {
                    lhs: Arc::new(self.clone()),
                    rhs: Arc::new(other.clone()),
                };
            }
            return Ok(result);
        }

        let mut result = self.elementwise_operation(other, |a, b| a * b)?;

        // Propagate requires_grad and record the operation for autograd
        if crate::should_record_grad(self.requires_grad || other.requires_grad) {
            result.requires_grad = true;
            result.operation = crate::Operation::Mul {
                lhs: Arc::new(self.clone()),
                rhs: Arc::new(other.clone()),
            };
        }

        Ok(result)
    }

    /// Element-wise division with another tensor (supports broadcasting).
    ///
    /// When either operand requires gradients the result records
    /// `Operation::Div` with the *un-broadcast* operands, so backward can apply
    /// `d/dlhs = 1 / rhs`, `d/drhs = -lhs / rhs²` and fold each gradient back to
    /// its operand's own shape.
    pub fn div(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Div<Output = T>,
    {
        // Hardware SIMD fast path (f32/f64, ≥ SIMD_BINARY_THRESHOLD elements),
        // reading both operands in place and writing one pooled output buffer.
        #[cfg(feature = "simd")]
        if let Some(mut result) = self.try_simd_binary(other, BinaryF32Op::Div)? {
            if crate::should_record_grad(self.requires_grad || other.requires_grad) {
                result.requires_grad = true;
                result.operation = crate::Operation::Div {
                    lhs: Arc::new(self.clone()),
                    rhs: Arc::new(other.clone()),
                };
            }
            return Ok(result);
        }

        let mut result = self.elementwise_operation(other, |a, b| a / b)?;

        // Propagate requires_grad and record the operation for autograd
        if crate::should_record_grad(self.requires_grad || other.requires_grad) {
            result.requires_grad = true;
            result.operation = crate::Operation::Div {
                lhs: Arc::new(self.clone()),
                rhs: Arc::new(other.clone()),
            };
        }

        Ok(result)
    }

    /// Handle broadcasting binary operations
    fn broadcast_binary_op<F>(&self, other: &Self, op: F) -> Result<Self>
    where
        F: Fn(T, T) -> T + Send + Sync,
    {
        use crate::broadcast::BroadcastOps;

        let self_shape_binding = self.shape();
        let self_shape = self_shape_binding.dims();
        let other_shape_binding = other.shape();
        let other_shape = other_shape_binding.dims();

        // Compute the broadcast result shape
        let broadcast_shape = BroadcastOps::compute_broadcast_shape(self_shape, other_shape)?;

        let total_elements = broadcast_shape.iter().product::<usize>();
        let result_data = self.with_operand_slices(other, |self_data, other_data| {
            let mut buf = global_acquire_uninit::<T>(total_elements);
            let uninit = &mut buf.as_uninit_slice_mut()[..total_elements];
            let mut count = 0;

            // Generate all possible indices for the broadcast shape
            let mut indices = vec![0; broadcast_shape.len()];
            for _ in 0..total_elements {
                // Map broadcast indices to original tensor indices
                let self_idx =
                    self.compute_broadcast_index(&indices, self_shape, &broadcast_shape)?;
                let other_idx =
                    other.compute_broadcast_index(&indices, other_shape, &broadcast_shape)?;

                let lhs = *self_data
                    .get(self_idx)
                    .ok_or_else(|| TorshError::IndexError {
                        index: self_idx,
                        size: self_data.len(),
                    })?;
                let rhs = *other_data
                    .get(other_idx)
                    .ok_or_else(|| TorshError::IndexError {
                        index: other_idx,
                        size: other_data.len(),
                    })?;
                uninit[count].write(op(lhs, rhs));
                count += 1;

                // Increment indices (like an odometer)
                Self::increment_indices(&mut indices, &broadcast_shape);
            }

            Ok(buf.into_vec(count))
        })?;
        Self::from_data(result_data, broadcast_shape, self.device)
    }

    /// Helper function to increment multi-dimensional indices
    fn increment_indices(indices: &mut [usize], shape: &[usize]) {
        for i in (0..indices.len()).rev() {
            indices[i] += 1;
            if indices[i] < shape[i] {
                break;
            }
            indices[i] = 0;
        }
    }

    /// Compute flat index from broadcast indices for this tensor
    fn compute_broadcast_index(
        &self,
        broadcast_indices: &[usize],
        original_shape: &[usize],
        broadcast_shape: &[usize],
    ) -> Result<usize> {
        let ndim_diff = broadcast_shape.len() - original_shape.len();
        let mut flat_index = 0;
        let mut stride = 1;

        for i in (0..original_shape.len()).rev() {
            let broadcast_idx = broadcast_indices[ndim_diff + i];
            let original_size = original_shape[i];

            // Handle broadcasting: if original size is 1, use index 0
            let actual_idx = if original_size == 1 { 0 } else { broadcast_idx };

            flat_index += actual_idx * stride;
            stride *= original_size;
        }

        Ok(flat_index)
    }

    /// Generic element-wise binary operation over two same-shape tensors.
    ///
    /// Both operands are read in place (no defensive copies) and the result is
    /// written into a single pooled buffer. Tensors with at least
    /// [`PARALLEL_ELEMENTWISE_THRESHOLD`] elements are split into
    /// cache-sized chunks across `scirs2_core::parallel_ops` workers; smaller
    /// ones run a single vectorisable pass.
    ///
    /// Element types with a hardware kernel (`f32`, `f64`) reach this function
    /// only below the SIMD threshold — see [`Tensor::try_simd_binary`].
    fn elementwise_operation<F>(&self, other: &Self, op: F) -> Result<Self>
    where
        F: Fn(T, T) -> T + Send + Sync,
    {
        // Handle broadcasting if shapes don't match
        if self.shape() != other.shape() {
            return self.broadcast_binary_op(other, op);
        }

        let n = self.numel();
        let device = self.device;
        let result_data = self.with_operand_slices(other, |a, b| {
            if a.len() != n || b.len() != n {
                return Err(TorshError::ShapeMismatch {
                    expected: vec![n],
                    got: vec![b.len()],
                });
            }

            let mut buf = global_acquire_uninit::<T>(n);
            {
                let uninit = &mut buf.as_uninit_slice_mut()[..n];

                // Operation-aware chunking: the chunk config that used to be
                // computed and discarded now actually sizes the work units.
                #[cfg(feature = "parallel")]
                {
                    if n >= PARALLEL_ELEMENTWISE_THRESHOLD {
                        use scirs2_core::parallel_ops::*;
                        let config = create_optimal_chunk_config(
                            n,
                            TensorOpType::ElementWise,
                            device,
                            false,
                        );
                        let chunk = config.max_chunk_size.max(config.min_chunk_size).max(1);
                        uninit
                            .par_chunks_mut(chunk)
                            .zip(a.par_chunks(chunk))
                            .zip(b.par_chunks(chunk))
                            .for_each(|((out_chunk, a_chunk), b_chunk)| {
                                for (slot, (&x, &y)) in
                                    out_chunk.iter_mut().zip(a_chunk.iter().zip(b_chunk.iter()))
                                {
                                    slot.write(op(x, y));
                                }
                            });
                        return Ok(buf.into_vec(n));
                    }
                }

                for (slot, (&x, &y)) in uninit.iter_mut().zip(a.iter().zip(b.iter())) {
                    slot.write(op(x, y));
                }
            }
            Ok(buf.into_vec(n))
        })?;

        Self::from_data(result_data, self.shape().dims().to_vec(), device)
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Block C – In-place tensor × tensor arithmetic with f32 SIMD fast paths
// ─────────────────────────────────────────────────────────────────────────────

impl<T: TensorElement + Copy> Tensor<T> {
    /// Element-wise in-place addition: `self += other`.
    ///
    /// Supports broadcasting `other` into `self`'s shape (e.g. adding a `[C]`
    /// bias to a `[B, C]` activation). If broadcasting would force `self` to
    /// grow, an error is returned — in-place ops cannot resize the destination.
    /// For f32 tensors with ≥ 1024 elements and matching shapes, this routes
    /// through the SIMD-backed `simd_ops_f32::add_assign_f32`. No new tensor
    /// data buffer is ever allocated.
    ///
    /// # Errors
    /// - `requires_grad` is true (autograd cannot be tracked through in-place mutations)
    /// - shapes are not broadcast-compatible
    /// - broadcasting would require `self` to grow
    pub fn add_(&mut self, other: &Self) -> Result<&mut Self>
    where
        T: std::ops::Add<Output = T>,
    {
        self.inplace_binary_op(
            other,
            "add_",
            crate::simd_ops_f32::add_assign_f32,
            |a, b| a + b,
        )
    }

    /// Element-wise in-place subtraction: `self -= other` (with broadcast of `other` into `self`).
    ///
    /// For f32 tensors with ≥ 1024 elements and matching shapes, this routes
    /// through `simd_ops_f32::sub_assign_f32`.
    pub fn sub_(&mut self, other: &Self) -> Result<&mut Self>
    where
        T: std::ops::Sub<Output = T>,
    {
        self.inplace_binary_op(
            other,
            "sub_",
            crate::simd_ops_f32::sub_assign_f32,
            |a, b| a - b,
        )
    }

    /// Element-wise in-place multiplication: `self *= other` (with broadcast of `other` into `self`).
    ///
    /// For f32 tensors with ≥ 1024 elements and matching shapes, this routes
    /// through `simd_ops_f32::mul_assign_f32`.
    pub fn mul_(&mut self, other: &Self) -> Result<&mut Self>
    where
        T: std::ops::Mul<Output = T>,
    {
        self.inplace_binary_op(
            other,
            "mul_",
            crate::simd_ops_f32::mul_assign_f32,
            |a, b| a * b,
        )
    }

    /// Element-wise in-place division: `self /= other` (with broadcast of `other` into `self`).
    ///
    /// For f32 tensors with ≥ 1024 elements and matching shapes, this routes
    /// through `simd_ops_f32::div_assign_f32`.
    pub fn div_(&mut self, other: &Self) -> Result<&mut Self>
    where
        T: std::ops::Div<Output = T>,
    {
        self.inplace_binary_op(
            other,
            "div_",
            crate::simd_ops_f32::div_assign_f32,
            |a, b| a / b,
        )
    }

    /// Shared in-place binary-op driver: validates shapes (with broadcast-into-self),
    /// then dispatches to the f32 SIMD fast path when applicable, otherwise to a
    /// single-locked slice loop. Never allocates a new data buffer.
    fn inplace_binary_op<S, F>(
        &mut self,
        other: &Self,
        op_name: &str,
        simd_f32: S,
        scalar_op: F,
    ) -> Result<&mut Self>
    where
        S: Fn(&mut [f32], &[f32]),
        F: Fn(T, T) -> T,
    {
        if self.requires_grad {
            return Err(TorshError::InvalidArgument(format!(
                "In-place operation `{op_name}` on tensor that requires grad is not allowed"
            )));
        }

        // f32 SIMD fast path: identical shapes, large enough to amortise dispatch.
        if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>()
            && self.shape() == other.shape()
            && self.numel() >= 1024
        {
            self.make_unique()?;
            let other_data = other.data()?;
            self.storage.with_slice_mut(|out_t: &mut [T]| {
                // Safety: TypeId guard above confirms T == f32, so reinterpreting
                // &mut [T] / &[T] as &mut [f32] / &[f32] is a no-op.
                let out_f32: &mut [f32] = unsafe {
                    std::slice::from_raw_parts_mut(out_t.as_mut_ptr() as *mut f32, out_t.len())
                };
                let rhs_f32: &[f32] = unsafe {
                    std::slice::from_raw_parts(other_data.as_ptr() as *const f32, other_data.len())
                };
                simd_f32(out_f32, rhs_f32);
                Ok(())
            })?;
            return Ok(self);
        }

        // Shape validation + broadcast-into-self path.
        let self_dims_vec = self.shape().dims().to_vec();
        let other_dims_vec = other.shape().dims().to_vec();
        let other_data = other.data()?;

        if self_dims_vec == other_dims_vec {
            // Equal shapes, generic-T scalar pass.
            self.make_unique()?;
            self.storage.with_slice_mut(|slice: &mut [T]| {
                debug_assert_eq!(slice.len(), other_data.len());
                for (dst, src) in slice.iter_mut().zip(other_data.iter()) {
                    *dst = scalar_op(*dst, *src);
                }
                Ok(())
            })?;
            return Ok(self);
        }

        // Broadcast path: must be broadcast-compatible AND self's shape must
        // already equal the broadcast shape (cannot grow self in place).
        use crate::broadcast::BroadcastOps;
        let broadcast_shape =
            BroadcastOps::compute_broadcast_shape(&self_dims_vec, &other_dims_vec)?;
        if broadcast_shape != self_dims_vec {
            return Err(TorshError::ShapeMismatch {
                expected: self_dims_vec,
                got: broadcast_shape,
            });
        }

        self.make_unique()?;
        let total = self_dims_vec.iter().product::<usize>();
        let self_dims = self_dims_vec.as_slice();
        let other_dims = other_dims_vec.as_slice();
        self.storage.with_slice_mut(|slice: &mut [T]| {
            for flat_idx in 0..total {
                let multi_index = BroadcastOps::flat_to_multi_index(flat_idx, self_dims);
                let other_idx =
                    BroadcastOps::compute_broadcast_index(&multi_index, other_dims, self_dims)?;
                slice[flat_idx] = scalar_op(slice[flat_idx], other_data[other_idx]);
            }
            Ok(())
        })?;

        Ok(self)
    }
}

// Internal operations for autograd (general implementations)
impl<T: TensorElement + Copy> Tensor<T> {
    /// Add operation (used by autograd backward pass)
    pub fn add_op(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Add<Output = T>,
    {
        self.add(other)
    }

    /// Multiply operation (used by autograd backward pass)
    pub fn mul_op(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Mul<Output = T>,
    {
        self.mul(other)
    }

    /// Sigmoid activation function with SIMD optimization
    pub fn sigmoid(&self) -> Result<Self>
    where
        T: torsh_core::dtype::FloatElement,
    {
        let result = self.sigmoid_forward()?;
        Ok(self.record_unary(result, UnaryKind::Sigmoid))
    }

    /// Forward-only sigmoid (all dispatch paths); autograd is recorded by the
    /// public [`Tensor::sigmoid`] wrapper.
    fn sigmoid_forward(&self) -> Result<Self>
    where
        T: torsh_core::dtype::FloatElement,
    {
        // GPU fast path: f32 CUDA tensors dispatch to oxicuda's ComputeBackend.
        // Declines to None (CPU fallback) unless a GPU backend is active.
        #[cfg(feature = "gpu")]
        if let Some(result) =
            crate::gpu_dispatch::try_unary_f32(self, crate::gpu_dispatch::UnaryOp::Sigmoid)
        {
            return Ok(result);
        }

        // ✅ SciRS2 SIMD Optimization - Vectorized sigmoid for f32 tensors
        #[cfg(feature = "simd")]
        {
            if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() > 1000 {
                return self.simd_sigmoid_f32();
            }
        }

        // ✅ SciRS2 Parallel Processing - Use parallel computation for medium tensors
        #[cfg(feature = "parallel")]
        {
            if self.numel() > 100 {
                let one = <T as scirs2_core::numeric::One>::one();
                return self.parallel_map(|x| {
                    // sigmoid(x) = 1 / (1 + exp(-x))
                    one / (one + (-x).exp())
                });
            }
        }

        // Fallback to sequential tensor operations
        let one = <T as scirs2_core::numeric::One>::one();
        let neg_self = self.neg()?;
        let exp_neg = neg_self.exp()?;
        let one_plus_exp = exp_neg.add_scalar(one)?;
        let ones = Self::ones(self.shape().dims(), self.device)?;
        ones.div(&one_plus_exp)
    }

    /// SIMD-optimized sigmoid activation function for f32 tensors
    #[cfg(feature = "simd")]
    fn simd_sigmoid_f32(&self) -> Result<Self> {
        self.simd_activation_f32(adaptive_simd::adaptive_simd_sigmoid_f32)
    }

    /// Shared driver for the f32 SIMD activations.
    ///
    /// Reads the operand straight out of storage instead of copying it with
    /// `data()`, and reinterprets the kernel's result buffer as `Vec<T>` in a
    /// single step rather than transmuting element by element. Only the one
    /// buffer produced by the kernel is allocated.
    ///
    /// # Safety contract
    /// Callers must have established `T == f32` (the `TypeId` guards in
    /// [`Tensor::relu`] and [`Tensor::sigmoid`] do so).
    #[cfg(feature = "simd")]
    fn simd_activation_f32<K>(&self, kernel: K) -> Result<Self>
    where
        K: Fn(&scirs2_core::ndarray::ArrayView1<f32>) -> Array1<f32>,
    {
        use scirs2_core::ndarray::ArrayView1;

        let numel = self.numel();
        let result_data: Vec<T> = self.with_contiguous_data(|data| {
            // SAFETY: the caller's TypeId guard confirms `T == f32`, so the
            // slice can be reinterpreted without copying.
            let data_f32: &[f32] =
                unsafe { std::slice::from_raw_parts(data.as_ptr().cast::<f32>(), numel) };
            let result_array = kernel(&ArrayView1::from(data_f32));
            let (values, _offset) = result_array.into_raw_vec_and_offset();
            // SAFETY: `T == f32`, so the buffer is already a valid `Vec<T>`;
            // reinterpreting it whole avoids a per-element `transmute_copy`
            // and a second allocation.
            Ok(unsafe {
                let mut values = std::mem::ManuallyDrop::new(values);
                Vec::from_raw_parts(
                    values.as_mut_ptr().cast::<T>(),
                    values.len(),
                    values.capacity(),
                )
            })
        })?;

        Self::from_data(result_data, self.shape().dims().to_vec(), self.device)
    }

    /// ReLU activation function (Rectified Linear Unit) with SIMD optimization
    pub fn relu(&self) -> Result<Self>
    where
        T: std::cmp::PartialOrd + scirs2_core::numeric::Zero,
    {
        let result = self.relu_forward()?;
        Ok(self.record_unary(result, UnaryKind::Relu))
    }

    /// Forward-only ReLU (all dispatch paths); autograd is recorded by the
    /// public [`Tensor::relu`] wrapper.
    fn relu_forward(&self) -> Result<Self>
    where
        T: std::cmp::PartialOrd + scirs2_core::numeric::Zero,
    {
        // GPU fast path: f32 CUDA tensors dispatch to oxicuda's ComputeBackend.
        // Declines to None (CPU fallback) unless a GPU backend is active.
        #[cfg(feature = "gpu")]
        if let Some(result) =
            crate::gpu_dispatch::try_unary_f32(self, crate::gpu_dispatch::UnaryOp::Relu)
        {
            return Ok(result);
        }

        let zero = <T as scirs2_core::numeric::Zero>::zero();

        // ✅ SciRS2 SIMD Optimization - Vectorized ReLU for f32 tensors
        #[cfg(feature = "simd")]
        {
            if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() > 1000 {
                return self.simd_relu_f32();
            }
        }

        // ✅ SciRS2 Parallel Processing - Use parallel map for medium tensors
        #[cfg(feature = "parallel")]
        {
            if self.numel() > 100 {
                return self.parallel_map(|x| if x > zero { x } else { zero });
            }
        }

        // Fallback to sequential processing
        self.map(|x| if x > zero { x } else { zero })
    }

    /// SIMD-optimized ReLU activation function for f32 tensors
    #[cfg(feature = "simd")]
    fn simd_relu_f32(&self) -> Result<Self> {
        self.simd_activation_f32(adaptive_simd::adaptive_simd_relu_f32)
    }

    /// SciRS2 Intelligent parallel map operation for medium-sized tensors with operation-aware chunking
    #[cfg(feature = "parallel")]
    fn parallel_map<F>(&self, op: F) -> Result<Self>
    where
        F: Fn(T) -> T + Send + Sync,
    {
        let data = self.data()?;

        // Use intelligent chunking system - activation functions get specialized chunking strategy
        let result_data = intelligent_parallel_process(
            data.iter().copied().collect::<Vec<_>>(),
            TensorOpType::Activation, // Most parallel_map calls are for activation functions
            self.device.clone(),
            op,
        );

        Self::from_data(result_data, self.shape().dims().to_vec(), self.device)
    }

    /// Element-wise minimum with another tensor (supports broadcasting).
    ///
    /// Records [`Operation::Minimum`] with the *un-broadcast* operands, so
    /// backward sends each element's gradient to whichever operand was smaller
    /// (splitting `0.5 / 0.5` on an exact tie) and folds it back to that
    /// operand's own shape.
    pub fn minimum(&self, other: &Self) -> Result<Self>
    where
        T: std::cmp::PartialOrd,
    {
        let result = self.elementwise_operation(other, |a, b| if a < b { a } else { b })?;
        Ok(self.record_extremum(other, result, false))
    }

    /// Element-wise maximum with another tensor (supports broadcasting).
    ///
    /// Records [`Operation::Maximum`] with the *un-broadcast* operands. This is
    /// what makes `relu(x) == x.maximum(&zeros)` differentiable, which is the
    /// form `torsh_nn::functional::relu` uses.
    pub fn maximum(&self, other: &Self) -> Result<Self>
    where
        T: std::cmp::PartialOrd,
    {
        let result = self.elementwise_operation(other, |a, b| if a > b { a } else { b })?;
        Ok(self.record_extremum(other, result, true))
    }

    /// Clamp tensor values between `min` and `max` bounds.
    ///
    /// Records [`Operation::ClampBounds`]: the gradient passes wherever the
    /// element was left alone (`min <= x <= max`, inclusive at both bounds, as
    /// in ATen's `clamp_backward`) and is zero wherever the clamp moved it.
    pub fn clamp(&self, min: T, max: T) -> Result<Self>
    where
        T: std::cmp::PartialOrd + Copy,
    {
        self.clamp_bounds(Some(min), Some(max))
    }

    /// Clamp tensor values from below only — PyTorch's `tensor.clamp_min(min)`.
    ///
    /// Gradient passes wherever `x >= min` (inclusive) and is zero below it.
    pub fn clamp_min(&self, min: T) -> Result<Self>
    where
        T: std::cmp::PartialOrd + Copy,
    {
        self.clamp_bounds(Some(min), None)
    }

    /// Clamp tensor values from above only — PyTorch's `tensor.clamp_max(max)`.
    ///
    /// Gradient passes wherever `x <= max` (inclusive) and is zero above it.
    pub fn clamp_max(&self, max: T) -> Result<Self>
    where
        T: std::cmp::PartialOrd + Copy,
    {
        self.clamp_bounds(None, Some(max))
    }

    /// Shared implementation of [`Tensor::clamp`], [`Tensor::clamp_min`] and
    /// [`Tensor::clamp_max`]: apply whichever bounds were supplied, then record
    /// them so backward can reproduce the same predicate.
    fn clamp_bounds(&self, min: Option<T>, max: Option<T>) -> Result<Self>
    where
        T: std::cmp::PartialOrd + Copy,
    {
        let data = self.to_vec()?;
        // Exactly the original `if x < min { min } else if x > max { max }
        // else { x }` ladder, with an absent bound skipping its arm. Order
        // matters and is preserved: for the degenerate `min > max` the lower
        // bound still wins, as it did before.
        let clamped_data: Vec<T> = data
            .iter()
            .map(|&x| match (min, max) {
                (Some(lower), _) if x < lower => lower,
                (_, Some(upper)) if x > upper => upper,
                _ => x,
            })
            .collect();

        let mut result = Self::from_data(
            clamped_data,
            self.shape().dims().to_vec(),
            self.device.clone(),
        )?;

        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::ClampBounds {
                input: Arc::new(self.clone()),
                min,
                max,
            };
        }

        Ok(result)
    }

    // ✅ clamp_ moved to in-place operations section below for PyTorch compatibility

    /// Dot product with another tensor (for 1D tensors)
    pub fn dot(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Mul<Output = T> + std::ops::Add<Output = T> + scirs2_core::numeric::Zero,
    {
        // For now, implement as element-wise multiply then sum
        let elementwise = self.mul(other)?;
        elementwise.sum()
    }
}

// SciRS2 backend integration for optimized operations
impl<T: TensorElement + Copy + scirs2_core::numeric::FromPrimitive> Tensor<T> {
    /// Use SciRS2 backend for optimized tensor addition
    pub fn add_scirs2(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Add<Output = T> + scirs2_core::numeric::Float,
    {
        // TODO: Integrate with actual SciRS2 backend
        // For now, fall back to basic implementation
        self.add(other)
    }

    /// Use SciRS2 backend for optimized tensor multiplication
    pub fn mul_scirs2(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Mul<Output = T> + scirs2_core::numeric::Float,
    {
        // TODO: Integrate with actual SciRS2 backend
        // For now, fall back to basic implementation
        self.mul(other)
    }

    /// Use SciRS2 backend for optimized tensor subtraction
    pub fn sub_scirs2(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Sub<Output = T> + scirs2_core::numeric::Float,
    {
        // TODO: Integrate with actual SciRS2 backend
        // For now, fall back to basic implementation
        self.sub(other)
    }

    /// Use SciRS2 backend for optimized tensor division
    pub fn div_scirs2(&self, other: &Self) -> Result<Self>
    where
        T: std::ops::Div<Output = T> + scirs2_core::numeric::Float,
    {
        // TODO: Integrate with actual SciRS2 backend
        // For now, fall back to basic implementation
        self.div(other)
    }
}

// Operator overloads for convenient syntax
impl<T: TensorElement + Copy> std::ops::Add for &Tensor<T>
where
    T: std::ops::Add<Output = T>,
{
    type Output = Tensor<T>;

    fn add(self, rhs: Self) -> Self::Output {
        self.add(rhs).expect("tensor addition should succeed")
    }
}

impl<T: TensorElement + Copy> std::ops::Sub for &Tensor<T>
where
    T: std::ops::Sub<Output = T>,
{
    type Output = Tensor<T>;

    fn sub(self, rhs: Self) -> Self::Output {
        self.sub(rhs).expect("tensor subtraction should succeed")
    }
}

impl<T: TensorElement + Copy> std::ops::Mul for &Tensor<T>
where
    T: std::ops::Mul<Output = T>,
{
    type Output = Tensor<T>;

    fn mul(self, rhs: Self) -> Self::Output {
        self.mul(rhs).expect("tensor multiplication should succeed")
    }
}

impl<T: TensorElement + Copy> std::ops::Div for &Tensor<T>
where
    T: std::ops::Div<Output = T>,
{
    type Output = Tensor<T>;

    fn div(self, rhs: Self) -> Self::Output {
        self.div(rhs).expect("tensor division should succeed")
    }
}

// Negation operator
impl<T: TensorElement + Copy> std::ops::Neg for &Tensor<T>
where
    T: std::ops::Neg<Output = T>,
{
    type Output = Tensor<T>;

    fn neg(self) -> Self::Output {
        let mut result = self.map(|x| -x).expect("negation map should succeed");
        // `map` is forward-only, so the operator has to record its own node;
        // negation is `MulScalar(-1)`, whose backward arm already exists.
        if crate::should_record_grad(self.requires_grad()) {
            result.requires_grad = true;
            result.operation = Operation::MulScalar {
                input: Arc::new(self.clone()),
                scalar: -<T as TensorElement>::one(),
            };
        }
        result
    }
}

// ✅ SciRS2 ADVANCED PERFORMANCE FEATURES
// High-performance implementations leveraging SciRS2 ecosystem

impl<T: TensorElement + Copy + scirs2_core::numeric::Float> Tensor<T> {
    /// Element-wise addition with SIMD acceleration (SciRS2)
    ///
    /// Uses real hardware SIMD instructions (AVX2/NEON) via `scirs2_core::simd_ops::SimdUnifiedOps`
    #[cfg(feature = "simd")]
    pub fn add_simd(&self, other: &Self) -> Result<Self>
    where
        T: scirs2_core::simd_ops::SimdUnifiedOps,
    {
        use scirs2_core::ndarray::Array1;

        // Shape check
        if self.shape().dims() != other.shape().dims() {
            return Err(TorshError::ShapeMismatch {
                expected: self.shape().dims().to_vec(),
                got: other.shape().dims().to_vec(),
            });
        }

        // Get data as vectors
        let data_a = self.to_vec()?;
        let data_b = other.to_vec()?;

        // Create ndarray arrays
        let arr_a = Array1::from_vec(data_a);
        let arr_b = Array1::from_vec(data_b);

        // Use REAL SIMD operation (not Rayon!)
        let result_arr = T::simd_add(&arr_a.view(), &arr_b.view());

        // Convert back to Tensor
        Tensor::from_vec(result_arr.to_vec(), self.shape().dims())
    }

    /// Element-wise multiplication with SIMD acceleration (SciRS2)
    ///
    /// Uses real hardware SIMD instructions (AVX2/NEON) via `scirs2_core::simd_ops::SimdUnifiedOps`
    #[cfg(feature = "simd")]
    pub fn mul_simd(&self, other: &Self) -> Result<Self>
    where
        T: scirs2_core::simd_ops::SimdUnifiedOps,
    {
        use scirs2_core::ndarray::Array1;

        // Shape check
        if self.shape().dims() != other.shape().dims() {
            return Err(TorshError::ShapeMismatch {
                expected: self.shape().dims().to_vec(),
                got: other.shape().dims().to_vec(),
            });
        }

        // Get data as vectors
        let data_a = self.to_vec()?;
        let data_b = other.to_vec()?;

        // Create ndarray arrays
        let arr_a = Array1::from_vec(data_a);
        let arr_b = Array1::from_vec(data_b);

        // Use REAL SIMD operation (not Rayon!)
        let result_arr = T::simd_mul(&arr_a.view(), &arr_b.view());

        // Convert back to Tensor
        Tensor::from_vec(result_arr.to_vec(), self.shape().dims())
    }

    /// Dot product with SIMD acceleration (SciRS2)
    ///
    /// Uses real hardware SIMD instructions (AVX2/NEON) via `scirs2_core::simd_ops::SimdUnifiedOps`
    ///
    /// Returns a scalar value (sum of element-wise products)
    #[cfg(feature = "simd")]
    pub fn dot_simd(&self, other: &Self) -> Result<T>
    where
        T: scirs2_core::simd_ops::SimdUnifiedOps,
    {
        use scirs2_core::ndarray::Array1;

        // Shape check
        if self.shape().dims() != other.shape().dims() {
            return Err(TorshError::ShapeMismatch {
                expected: self.shape().dims().to_vec(),
                got: other.shape().dims().to_vec(),
            });
        }

        // Get data as vectors
        let data_a = self.to_vec()?;
        let data_b = other.to_vec()?;

        // Create ndarray arrays
        let arr_a = Array1::from_vec(data_a);
        let arr_b = Array1::from_vec(data_b);

        // Use REAL SIMD dot product (not Rayon!)
        Ok(T::simd_dot(&arr_a.view(), &arr_b.view()))
    }

    /// Memory-efficient reduction using SciRS2 intelligent chunking and lazy evaluation
    pub fn reduce_memory_efficient<F>(&self, func: F) -> Result<T>
    where
        F: Fn(T, T) -> T + Send + Sync,
    {
        #[cfg(feature = "profiling")]
        {
            // let _profile = profile_section!("tensor_reduce_memory_efficient");
        }

        // Use simple reduction for now to get basic functionality working
        // Regular reduction for smaller tensors
        let data = self.to_vec()?;
        Ok(data
            .into_iter()
            .reduce(func)
            .unwrap_or_else(|| <T as scirs2_core::numeric::Zero>::zero()))
    }
}

// ✅ In-place activation functions for PyTorch compatibility
impl<T: TensorElement + Copy + std::ops::Mul<Output = T>> Tensor<T> {
    /// In-place ReLU activation: self = max(0, self)
    ///
    /// # PyTorch Compatibility
    /// Equivalent to PyTorch's `tensor.relu_()`
    ///
    /// # Errors
    /// - Returns error if `requires_grad` is true
    pub fn relu_(&mut self) -> Result<&mut Self>
    where
        T: std::cmp::PartialOrd + scirs2_core::numeric::Zero,
    {
        if self.requires_grad {
            return Err(TorshError::InvalidArgument(
                "In-place operation on tensor that requires grad is not allowed".to_string(),
            ));
        }

        // f32 SIMD fast path with PyTorch NaN-passthrough semantics
        if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() >= 1024 {
            self.make_unique()?;
            self.storage.with_slice_mut(|out_t: &mut [T]| {
                // Safety: TypeId confirmed T == f32.
                let out_f32: &mut [f32] = unsafe {
                    std::slice::from_raw_parts_mut(out_t.as_mut_ptr() as *mut f32, out_t.len())
                };
                crate::simd_ops_f32::relu_assign_f32(out_f32);
                Ok(())
            })?;
            return Ok(self);
        }

        // Generic path: one copy-on-write promotion, then a single locked pass
        // over the buffer (`apply_`) instead of a lock pair per element.
        let zero = <T as scirs2_core::numeric::Zero>::zero();
        self.apply_(|x| if x < zero { zero } else { x })?;

        Ok(self)
    }

    /// In-place sigmoid activation: self = 1 / (1 + exp(-self))
    ///
    /// # PyTorch Compatibility
    /// Equivalent to PyTorch's `tensor.sigmoid_()`
    pub fn sigmoid_(&mut self) -> Result<&mut Self>
    where
        T: torsh_core::dtype::FloatElement,
    {
        if self.requires_grad {
            return Err(TorshError::InvalidArgument(
                "In-place operation on tensor that requires grad is not allowed".to_string(),
            ));
        }

        // One copy-on-write promotion, then a single locked pass.
        let one = <T as scirs2_core::numeric::One>::one();
        self.apply_(|x| one / (one + (-x).exp()))?;

        Ok(self)
    }

    /// In-place tanh activation: self = tanh(self)
    ///
    /// # PyTorch Compatibility
    /// Equivalent to PyTorch's `tensor.tanh_()`
    pub fn tanh_(&mut self) -> Result<&mut Self>
    where
        T: torsh_core::dtype::FloatElement,
    {
        if self.requires_grad {
            return Err(TorshError::InvalidArgument(
                "In-place operation on tensor that requires grad is not allowed".to_string(),
            ));
        }

        // One copy-on-write promotion, then a single locked pass.
        self.apply_(|x| x.tanh())?;

        Ok(self)
    }

    /// In-place GELU activation
    ///
    /// # PyTorch Compatibility
    /// Equivalent to PyTorch's `tensor.gelu_()`
    pub fn gelu_(&mut self) -> Result<&mut Self>
    where
        T: torsh_core::dtype::FloatElement,
    {
        if self.requires_grad {
            return Err(TorshError::InvalidArgument(
                "In-place operation on tensor that requires grad is not allowed".to_string(),
            ));
        }

        // Constants of the tanh GELU approximation. `T::from` is fallible for
        // exotic element types, so it is reported rather than unwrapped.
        let constant = |value: f64| -> Result<T> {
            T::from(value).ok_or_else(|| {
                TorshError::InvalidArgument(format!(
                    "gelu_: element type cannot represent the constant {value}"
                ))
            })
        };
        let pi = constant(std::f64::consts::PI)?;
        let two = constant(2.0)?;
        let sqrt_2_over_pi = (two / pi).sqrt();
        let point_044715 = constant(0.044715)?;
        let one = <T as scirs2_core::numeric::One>::one();
        let half = constant(0.5)?;

        // One copy-on-write promotion, then a single locked pass.
        self.apply_(|x| {
            let x_cubed = x * x * x;
            let tanh_input = sqrt_2_over_pi * (x + point_044715 * x_cubed);
            half * x * (one + tanh_input.tanh())
        })?;

        Ok(self)
    }

    /// In-place leaky ReLU activation
    ///
    /// # PyTorch Compatibility
    /// Equivalent to PyTorch's `tensor.leaky_relu_(negative_slope)`
    pub fn leaky_relu_(&mut self, negative_slope: T) -> Result<&mut Self>
    where
        T: std::cmp::PartialOrd + scirs2_core::numeric::Zero,
    {
        if self.requires_grad {
            return Err(TorshError::InvalidArgument(
                "In-place operation on tensor that requires grad is not allowed".to_string(),
            ));
        }

        // f32 SIMD fast path with PyTorch NaN-passthrough semantics
        if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() >= 1024 {
            // Safety: TypeId confirmed T == f32; copy the scalar via transmute_copy.
            let slope_f32: f32 = unsafe { std::mem::transmute_copy::<T, f32>(&negative_slope) };
            self.make_unique()?;
            self.storage.with_slice_mut(|out_t: &mut [T]| {
                let out_f32: &mut [f32] = unsafe {
                    std::slice::from_raw_parts_mut(out_t.as_mut_ptr() as *mut f32, out_t.len())
                };
                crate::simd_ops_f32::leaky_relu_assign_f32(out_f32, slope_f32);
                Ok(())
            })?;
            return Ok(self);
        }

        // One copy-on-write promotion, then a single locked pass.
        let zero = <T as scirs2_core::numeric::Zero>::zero();
        self.apply_(|x| if x < zero { negative_slope * x } else { x })?;

        Ok(self)
    }

    /// In-place clamp operation: self = clamp(self, min, max)
    ///
    /// # PyTorch Compatibility
    /// Equivalent to PyTorch's `tensor.clamp_(min, max)`
    pub fn clamp_(&mut self, min: T, max: T) -> Result<&mut Self>
    where
        T: std::cmp::PartialOrd,
    {
        if self.requires_grad {
            return Err(TorshError::InvalidArgument(
                "In-place operation on tensor that requires grad is not allowed".to_string(),
            ));
        }

        // f32 SIMD fast path with PyTorch NaN-passthrough semantics
        if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() >= 1024 {
            // Safety: TypeId confirmed T == f32; copy scalars via transmute_copy.
            let min_f32: f32 = unsafe { std::mem::transmute_copy::<T, f32>(&min) };
            let max_f32: f32 = unsafe { std::mem::transmute_copy::<T, f32>(&max) };
            self.make_unique()?;
            self.storage.with_slice_mut(|out_t: &mut [T]| {
                let out_f32: &mut [f32] = unsafe {
                    std::slice::from_raw_parts_mut(out_t.as_mut_ptr() as *mut f32, out_t.len())
                };
                crate::simd_ops_f32::clamp_assign_f32(out_f32, min_f32, max_f32);
                Ok(())
            })?;
            return Ok(self);
        }

        // One copy-on-write promotion, then a single locked pass.
        self.apply_(|x| {
            if x < min {
                min
            } else if x > max {
                max
            } else {
                x
            }
        })?;

        Ok(self)
    }
}

#[path = "math_ops_trig.rs"]
mod math_ops_trig;

#[cfg(test)]
#[path = "math_ops_tests.rs"]
mod tests;