torsh_tensor/

shape_ops.rs

//! Shape and view operations for tensors
//!
//! This module provides comprehensive tensor shape manipulation and view operations
//! including reshaping, transposing, slicing, squeezing, unsqueezing, and permuting.
//!
//! # Features
//!
//! - **Zero-copy views**: Efficient view operations that share underlying data
//! - **Safe reshaping**: Comprehensive validation and overflow checking
//! - **Dimension manipulation**: Squeeze, unsqueeze, transpose, and permute operations
//! - **Slicing operations**: Flexible tensor slicing with stride computation
//! - **Broadcasting support**: Expand operations for broadcasting compatibility
//! - **Contiguity checking**: Efficient memory layout validation

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

use crate::core_ops::{Operation, Tensor};
use crate::memory_pool::global_acquire_uninit;

impl<T: TensorElement + Copy> Tensor<T> {
    /// Get size of a specific dimension
    pub fn size(&self, dim: i32) -> Result<usize> {
        self.shape().size(dim)
    }

    /// Reshapes the tensor to a new shape (creates a view or copy if needed).
    ///
    /// This is equivalent to PyTorch's `view()` operation. The total number of elements
    /// must remain the same. You can use `-1` for one dimension to have it inferred automatically.
    ///
    /// # Arguments
    ///
    /// * `shape` - The new shape as a slice of dimensions. Use `-1` to infer one dimension.
    ///
    /// # Returns
    ///
    /// A reshaped tensor, or an error if the reshape is invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// use torsh_tensor::creation::zeros;
    ///
    /// // Reshape a 1D tensor to 2D
    /// let t = zeros::<f32>(&[6]).expect("tensor creation should succeed");
    /// let reshaped = t.view(&[2, 3]).expect("view should succeed");
    /// assert_eq!(reshaped.shape().dims(), &[2, 3]);
    ///
    /// // Use -1 to infer a dimension
    /// let t2 = zeros::<f32>(&[12]).expect("tensor creation should succeed");
    /// let auto = t2.view(&[-1, 4]).expect("view should succeed");  // Infers 3 for first dimension
    /// assert_eq!(auto.shape().dims(), &[3, 4]);
    ///
    /// // Flatten to 1D
    /// let matrix = zeros::<f32>(&[3, 4, 5]).expect("tensor creation should succeed");
    /// let flat = matrix.view(&[-1]).expect("view should succeed");
    /// assert_eq!(flat.shape().dims(), &[60]);
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - More than one dimension is `-1`
    /// - The total number of elements doesn't match
    /// - Any dimension would overflow
    ///
    /// # See Also
    ///
    /// * [`Self::reshape`] - Alias for `view()`
    /// * [`Self::view_as`] - Zero-copy view for compatible shapes
    /// * [`Self::contiguous`] - Make tensor contiguous in memory
    pub fn view(&self, shape: &[i32]) -> Result<Self> {
        // Validate that there's at most one -1 in the shape
        let infer_count = shape.iter().filter(|&&x| x == -1).count();
        if infer_count > 1 {
            return Err(TorshError::InvalidShape(
                "Only one dimension can be inferred (only one -1 allowed)".to_string(),
            ));
        }

        let new_shape: Result<Vec<usize>> = shape
            .iter()
            .map(|&d| {
                if d == -1 {
                    // Infer dimension - first validate all other dimensions are valid
                    let known_dims: Result<Vec<usize>> = shape
                        .iter()
                        .filter(|&&x| x != -1)
                        .map(|&x| {
                            if x < 0 {
                                Err(TorshError::InvalidShape(format!(
                                    "Invalid dimension size: {x} (negative dimensions not allowed except -1)"
                                )))
                            } else {
                                Ok(x as usize)
                            }
                        })
                        .collect();

                    let known_dims = known_dims?;

                    // Check for overflow in product calculation
                    let known_product = known_dims.iter().try_fold(1usize, |acc, &dim| {
                        acc.checked_mul(dim).ok_or_else(|| {
                            TorshError::InvalidShape(
                                "Shape dimensions too large (would overflow)".to_string()
                            )
                        })
                    })?;

                    if known_product == 0 {
                        return Err(TorshError::InvalidShape(
                            "Cannot infer dimension with zero-sized dimensions".to_string(),
                        ));
                    }

                    let total = self.numel();
                    if total % known_product != 0 {
                        return Err(TorshError::InvalidShape(
                            "Cannot infer dimension: size is not divisible".to_string(),
                        ));
                    }

                    Ok(total / known_product)
                } else if d < 0 {
                    Err(TorshError::InvalidShape(format!(
                        "Invalid dimension size: {d}"
                    )))
                } else {
                    Ok(d as usize)
                }
            })
            .collect();

        let new_shape = new_shape?;

        // Check for overflow in total elements calculation
        let new_numel = new_shape.iter().try_fold(1usize, |acc, &dim| {
            acc.checked_mul(dim).ok_or_else(|| {
                TorshError::InvalidShape(
                    "Reshaped tensor would be too large (would overflow)".to_string(),
                )
            })
        })?;

        if new_numel != self.numel() {
            return Err(TorshError::InvalidShape(format!(
                "Shape {:?} is invalid for tensor of size {}",
                new_shape,
                self.numel()
            )));
        }

        // Create a new tensor with the same data but different shape
        let data = self.to_vec()?;
        Self::from_data(data, new_shape, self.device)
    }

    /// Create an efficient view with different shape (shares data, no copying)
    /// This is the zero-copy version of view() for compatible shapes
    pub fn view_as(&self, shape: &[usize]) -> Result<Self> {
        // Validate that the total number of elements is the same
        let new_numel = shape.iter().product::<usize>();
        if new_numel != self.numel() {
            return Err(TorshError::InvalidShape(format!(
                "Shape {:?} is invalid for tensor of size {}",
                shape,
                self.numel()
            )));
        }

        // Only create efficient views for contiguous tensors or existing views
        // that are still relatively simple
        if !self.is_contiguous() {
            return Err(TorshError::InvalidShape(
                "Cannot create efficient view of non-contiguous tensor".to_string(),
            ));
        }

        // Create new tensor sharing the same storage
        Ok(Self {
            storage: self.storage.clone(),
            shape: Shape::new(shape.to_vec()),
            device: self.device,
            requires_grad: self.requires_grad,
            grad: Arc::new(RwLock::new(None)), // Views don't share gradients
            operation: Operation::Leaf,        // Views reset operation tracking
            strides: None,                     // Use default contiguous strides for simple reshapes
            storage_offset: self.storage_offset,
            base_tensor: if self.is_view() {
                // If this is already a view, keep reference to the original base
                self.base_tensor.clone()
            } else {
                // This is a base tensor, so create a weak reference to it
                Some(Arc::downgrade(&Arc::new(self.clone())))
            },
        })
    }

    /// Create a view of a slice along a dimension (shares data, no copying)
    pub fn slice_tensor(&self, dim: usize, start: usize, end: usize) -> Result<Self> {
        if dim >= self.ndim() {
            return Err(TorshError::InvalidArgument(format!(
                "Dimension {} out of range for tensor with {} dimensions",
                dim,
                self.ndim()
            )));
        }

        let shape = self.shape.dims();
        if start >= shape[dim] || end > shape[dim] || start >= end {
            return Err(TorshError::InvalidArgument(format!(
                "Invalid slice range [{}:{}] for dimension {} of size {}",
                start, end, dim, shape[dim]
            )));
        }

        // Calculate new shape
        let mut new_shape = shape.to_vec();
        new_shape[dim] = end - start;

        // Calculate new strides and offset
        let current_strides = self.strides();
        let offset_adjustment = start * current_strides[dim];

        Ok(Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: self.requires_grad,
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(current_strides),
            storage_offset: self.storage_offset + offset_adjustment,
            base_tensor: if self.is_view() {
                self.base_tensor.clone()
            } else {
                Some(Arc::downgrade(&Arc::new(self.clone())))
            },
        })
    }

    /// Create a transposed view (shares data, no copying)
    pub fn transpose_view(&self, dim0: usize, dim1: usize) -> Result<Self> {
        if dim0 >= self.ndim() || dim1 >= self.ndim() {
            return Err(TorshError::InvalidArgument(format!(
                "Dimensions {} and {} out of range for tensor with {} dimensions",
                dim0,
                dim1,
                self.ndim()
            )));
        }

        if dim0 == dim1 {
            return Ok(self.clone());
        }

        // Create new shape and strides
        let mut new_shape = self.shape.dims().to_vec();
        let mut new_strides = self.strides();

        // Swap dimensions
        new_shape.swap(dim0, dim1);
        new_strides.swap(dim0, dim1);

        Ok(Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: self.requires_grad,
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: if self.is_view() {
                self.base_tensor.clone()
            } else {
                Some(Arc::downgrade(&Arc::new(self.clone())))
            },
        })
    }

    /// Squeeze a tensor along a specific dimension (removes dimension of size 1)
    pub fn squeeze_tensor(&self, dim: usize) -> Result<Self> {
        if dim >= self.ndim() {
            return Err(TorshError::InvalidArgument(format!(
                "Dimension {} out of range for tensor with {} dimensions",
                dim,
                self.ndim()
            )));
        }

        let shape = self.shape.dims();
        if shape[dim] != 1 {
            return Err(TorshError::InvalidArgument(format!(
                "Cannot squeeze dimension {} of size {}",
                dim, shape[dim]
            )));
        }

        // Remove the dimension from shape and strides
        let mut new_shape = shape.to_vec();
        new_shape.remove(dim);

        let mut new_strides = self.strides();
        new_strides.remove(dim);

        Ok(Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: self.requires_grad,
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: if self.is_view() {
                self.base_tensor.clone()
            } else {
                Some(Arc::downgrade(&Arc::new(self.clone())))
            },
        })
    }

    /// Unsqueeze a tensor at a specific dimension (adds dimension of size 1)
    pub fn unsqueeze_tensor(&self, dim: usize) -> Result<Self> {
        if dim > self.ndim() {
            return Err(TorshError::InvalidArgument(format!(
                "Dimension {} out of range for insertion in tensor with {} dimensions",
                dim,
                self.ndim()
            )));
        }

        // Insert new dimension into shape and strides
        let mut new_shape = self.shape.dims().to_vec();
        new_shape.insert(dim, 1);

        let mut new_strides = self.strides();
        // For the new dimension, stride should be the product of all dimensions to the right
        let new_stride = if dim == new_shape.len() - 1 {
            1 // Last dimension always has stride 1
        } else {
            new_strides[dim] // Use the stride that was at this position
        };
        new_strides.insert(dim, new_stride);

        Ok(Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: self.requires_grad,
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: if self.is_view() {
                self.base_tensor.clone()
            } else {
                Some(Arc::downgrade(&Arc::new(self.clone())))
            },
        })
    }

    /// Transposes two dimensions of the tensor.
    ///
    /// Swaps the specified dimensions, creating a new tensor. For 2D tensors, calling
    /// `transpose(0, 1)` produces the standard matrix transpose operation.
    ///
    /// # Arguments
    ///
    /// * `dim0` - The first dimension to swap. Negative values count from the end.
    /// * `dim1` - The second dimension to swap. Negative values count from the end.
    ///
    /// # Returns
    ///
    /// A tensor with the specified dimensions transposed.
    ///
    /// # Examples
    ///
    /// ```
    /// use torsh_tensor::creation::{zeros, arange};
    ///
    /// // Standard matrix transpose
    /// let matrix = zeros::<f32>(&[3, 4]).expect("tensor creation should succeed");
    /// let transposed = matrix.transpose(0, 1).expect("transpose should succeed");
    /// assert_eq!(transposed.shape().dims(), &[4, 3]);
    ///
    /// // Transpose in 3D tensor
    /// let cube = zeros::<f32>(&[2, 3, 4]).expect("tensor creation should succeed");
    /// let swapped = cube.transpose(0, 2).expect("transpose should succeed");
    /// assert_eq!(swapped.shape().dims(), &[4, 3, 2]);
    ///
    /// // Use negative indexing
    /// let t = zeros::<f32>(&[5, 6, 7]).expect("tensor creation should succeed");
    /// let result = t.transpose(-2, -1).expect("transpose should succeed");
    /// assert_eq!(result.shape().dims(), &[5, 7, 6]);
    ///
    /// // Practical use: convert between row-major and column-major
    /// let data = arange(0, 12, 1).expect("arange should succeed");
    /// let row_major = data.reshape(&[3, 4]).expect("reshape should succeed");
    /// let col_major = row_major.transpose(0, 1).expect("transpose should succeed");
    /// ```
    ///
    /// # See Also
    ///
    /// * [`Self::permute`] - Rearrange dimensions in arbitrary order
    /// * [`Self::view`] - Reshape to different dimensions
    pub fn transpose(&self, dim0: i32, dim1: i32) -> Result<Self> {
        let ndim = self.ndim();
        let dim0 = if dim0 < 0 {
            (ndim as i32 + dim0) as usize
        } else {
            dim0 as usize
        };
        let dim1 = if dim1 < 0 {
            (ndim as i32 + dim1) as usize
        } else {
            dim1 as usize
        };

        if dim0 >= ndim || dim1 >= ndim {
            return Err(TorshError::InvalidArgument(format!(
                "Dimensions {} and {} out of range for tensor with {} dimensions",
                dim0, dim1, ndim
            )));
        }

        if ndim == 2 && dim0 != dim1 {
            self.transpose_2d()
        } else {
            self.transpose_view(dim0, dim1)
        }
    }

    /// 2D transpose implementation
    fn transpose_2d(&self) -> Result<Self> {
        let shape = self.shape.dims();
        if shape.len() != 2 {
            return Err(TorshError::InvalidArgument(
                "transpose_2d only works with 2D tensors".to_string(),
            ));
        }

        let (rows, cols) = (shape[0], shape[1]);
        let data = self.to_vec()?;
        let n = data.len();
        let mut buf = global_acquire_uninit::<T>(n);
        let uninit = buf.as_uninit_slice_mut();
        let mut count = 0;

        for col in 0..cols {
            for row in 0..rows {
                uninit[count].write(data[row * cols + col]);
                count += 1;
            }
        }

        let transposed_data = buf.into_vec(count);
        Self::from_data(transposed_data, vec![cols, rows], self.device)
    }

    /// Permute dimensions according to the given order
    pub fn permute(&self, dims: &[i32]) -> Result<Self> {
        let ndim = self.ndim();

        if dims.len() != ndim {
            return Err(TorshError::InvalidArgument(format!(
                "Number of dimensions in permutation ({}) doesn't match tensor dimensions ({})",
                dims.len(),
                ndim
            )));
        }

        // Convert negative indices and validate
        let perm_dims: Result<Vec<usize>> = dims
            .iter()
            .map(|&d| {
                let dim = if d < 0 { ndim as i32 + d } else { d } as usize;
                if dim >= ndim {
                    Err(TorshError::InvalidArgument(format!(
                        "Dimension {} out of range for tensor with {} dimensions",
                        d, ndim
                    )))
                } else {
                    Ok(dim)
                }
            })
            .collect();

        let perm_dims = perm_dims?;

        // Check for duplicates
        let mut sorted_dims = perm_dims.clone();
        sorted_dims.sort_unstable();
        for i in 0..ndim {
            if sorted_dims[i] != i {
                return Err(TorshError::InvalidArgument(
                    "Permutation must contain each dimension exactly once".to_string(),
                ));
            }
        }

        // Create new shape and strides
        let old_shape = self.shape.dims();
        let old_strides = self.strides();

        let new_shape: Vec<usize> = perm_dims.iter().map(|&i| old_shape[i]).collect();
        let new_strides: Vec<usize> = perm_dims.iter().map(|&i| old_strides[i]).collect();

        Ok(Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: self.requires_grad,
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: if self.is_view() {
                self.base_tensor.clone()
            } else {
                Some(Arc::downgrade(&Arc::new(self.clone())))
            },
        })
    }

    /// Removes a dimension of size 1 at the specified position.
    ///
    /// This operation reduces the dimensionality of the tensor by removing dimensions
    /// that have size 1. Commonly used to remove singleton dimensions after reductions
    /// or to match tensor shapes for operations.
    ///
    /// # Arguments
    ///
    /// * `dim` - The dimension to squeeze. Negative values count from the end.
    ///
    /// # Returns
    ///
    /// A tensor with the specified dimension removed, or an error if the dimension
    /// doesn't have size 1.
    ///
    /// # Examples
    ///
    /// ```
    /// use torsh_tensor::creation::zeros;
    ///
    /// // Remove a singleton dimension
    /// let t = zeros::<f32>(&[3, 1, 4]).expect("tensor creation should succeed");
    /// let squeezed = t.squeeze(1).expect("squeeze should succeed");
    /// assert_eq!(squeezed.shape().dims(), &[3, 4]);
    ///
    /// // Use negative indexing
    /// let t2 = zeros::<f32>(&[2, 3, 1]).expect("tensor creation should succeed");
    /// let squeezed2 = t2.squeeze(-1).expect("squeeze should succeed");
    /// assert_eq!(squeezed2.shape().dims(), &[2, 3]);
    ///
    /// // After a reduction operation
    /// let matrix = zeros::<f32>(&[5, 10]).expect("tensor creation should succeed");
    /// let reduced = matrix.sum_dim(&[1], true).expect("sum_dim should succeed");  // Shape: [5, 1]
    /// let final_result = reduced.squeeze(1).expect("squeeze should succeed");  // Shape: [5]
    /// ```
    ///
    /// # See Also
    ///
    /// * [`Self::squeeze_all`] - Remove all singleton dimensions
    /// * [`Self::unsqueeze`] - Add a singleton dimension
    pub fn squeeze(&self, dim: i32) -> Result<Self> {
        let ndim = self.ndim();
        let dim = if dim < 0 {
            (ndim as i32 + dim) as usize
        } else {
            dim as usize
        };

        self.squeeze_tensor(dim)
    }

    /// Squeeze all dimensions with size 1
    pub fn squeeze_all(&self) -> Result<Self> {
        let shape = self.shape.dims();
        let new_shape: Vec<usize> = shape.iter().copied().filter(|&s| s != 1).collect();

        if new_shape.is_empty() {
            // If all dimensions were 1, result should be a scalar (0-dimensional tensor)
            let data = self.to_vec()?;
            Self::from_data(data, vec![], self.device)
        } else {
            let data = self.to_vec()?;
            Self::from_data(data, new_shape, self.device)
        }
    }

    /// Adds a dimension of size 1 at the specified position.
    ///
    /// This operation increases the dimensionality of the tensor by inserting a new
    /// dimension of size 1. Commonly used to add batch dimensions or to match tensor
    /// shapes for broadcasting operations.
    ///
    /// # Arguments
    ///
    /// * `dim` - The position to insert the new dimension. Negative values count from the end.
    ///
    /// # Returns
    ///
    /// A tensor with an additional dimension of size 1 inserted.
    ///
    /// # Examples
    ///
    /// ```
    /// use torsh_tensor::creation::zeros;
    ///
    /// // Add a batch dimension at the beginning
    /// let t = zeros::<f32>(&[3, 4]).expect("tensor creation should succeed");
    /// let batched = t.unsqueeze(0).expect("unsqueeze should succeed");
    /// assert_eq!(batched.shape().dims(), &[1, 3, 4]);
    ///
    /// // Add a dimension at the end
    /// let t2 = zeros::<f32>(&[5]).expect("tensor creation should succeed");
    /// let expanded = t2.unsqueeze(-1).expect("unsqueeze should succeed");
    /// assert_eq!(expanded.shape().dims(), &[5, 1]);
    ///
    /// // Prepare for broadcasting
    /// let weights = zeros::<f32>(&[64]).expect("tensor creation should succeed");
    /// let weights_2d = weights.unsqueeze(0).expect("unsqueeze should succeed");  // Shape: [1, 64]
    /// // Now can broadcast with shape [batch_size, 64]
    /// ```
    ///
    /// # See Also
    ///
    /// * [`Self::squeeze`] - Remove a singleton dimension
    /// * [`Self::view`] - Reshape to arbitrary shape
    pub fn unsqueeze(&self, dim: i32) -> Result<Self> {
        let ndim = self.ndim();
        let dim = if dim < 0 {
            (ndim as i32 + dim + 1) as usize
        } else {
            dim as usize
        };

        self.unsqueeze_tensor(dim)
    }

    /// Reshapes the tensor to a new shape.
    ///
    /// This is an alias for [`view()`](Self::view) and provides the same functionality.
    /// The total number of elements must remain the same.
    ///
    /// # Arguments
    ///
    /// * `shape` - The new shape as a slice of dimensions. Use `-1` to infer one dimension.
    ///
    /// # Returns
    ///
    /// A reshaped tensor, or an error if the reshape is invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// use torsh_tensor::creation::arange;
    ///
    /// // Reshape a sequence to a matrix
    /// let t = arange(0, 12, 1).expect("arange should succeed");
    /// let matrix = t.reshape(&[3, 4]).expect("reshape should succeed");
    /// assert_eq!(matrix.shape().dims(), &[3, 4]);
    ///
    /// // Reshape with automatic dimension inference
    /// let cube = t.reshape(&[2, -1, 3]).expect("reshape should succeed");  // Infers 2 for middle dimension
    /// assert_eq!(cube.shape().dims(), &[2, 2, 3]);
    /// ```
    ///
    /// # See Also
    ///
    /// * [`Self::view`] - The underlying implementation
    pub fn reshape(&self, shape: &[i32]) -> Result<Self> {
        self.view(shape)
    }

    /// Check if tensor is contiguous in memory
    pub fn is_contiguous(&self) -> bool {
        // A tensor is contiguous if its strides match the default strides for its shape
        let default_strides = self.compute_default_strides();
        let current_strides = self.strides();

        current_strides == default_strides
    }

    /// Make tensor contiguous if it isn't already
    pub fn contiguous(&self) -> Result<Self> {
        if self.is_contiguous() {
            Ok(self.clone())
        } else {
            // Need to copy data to make it contiguous
            let data = self.to_vec()?;
            Self::from_data(data, self.shape.dims().to_vec(), self.device)
        }
    }

    /// Expand tensor to a larger size
    pub fn expand(&self, shape: &[usize]) -> Result<Self> {
        let old_shape = self.shape.dims();

        // Validate that expansion is possible
        if shape.len() < old_shape.len() {
            return Err(TorshError::InvalidShape(
                "Cannot expand to smaller number of dimensions".to_string(),
            ));
        }

        // Check dimension compatibility (broadcasting rules) and build the strides
        // for a zero-copy broadcast view.
        //
        // A dimension of size 1 that is expanded to size N is given stride 0, so all
        // N logical indices resolve to the same storage element (PyTorch `expand`
        // semantics). Newly-introduced leading dimensions are likewise stride 0, and
        // dimensions that keep their size retain their original (possibly
        // non-contiguous) stride. No tensor data is copied: the result shares the
        // source storage via `Arc`, so an expanded tensor never duplicates data.
        let offset = shape.len() - old_shape.len();
        let current_strides = self.strides();
        let mut new_strides = vec![0usize; shape.len()];
        for (i, &old_dim) in old_shape.iter().enumerate() {
            let new_dim = shape[offset + i];
            if old_dim == new_dim {
                new_strides[offset + i] = current_strides[i];
            } else if old_dim == 1 {
                new_strides[offset + i] = 0;
            } else {
                return Err(TorshError::InvalidShape(format!(
                    "Cannot expand dimension {} from {} to {}",
                    i, old_dim, new_dim
                )));
            }
        }
        // Any leading broadcast dimensions already have stride 0 from initialization.

        Ok(Self {
            storage: self.storage.clone(),
            shape: Shape::new(shape.to_vec()),
            device: self.device,
            requires_grad: false,
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: if self.is_view() {
                self.base_tensor.clone()
            } else {
                Some(Arc::downgrade(&Arc::new(self.clone())))
            },
        })
    }

    /// Move dimensions from source positions to destination positions
    ///
    /// # PyTorch Compatibility
    /// Equivalent to `torch.movedim(tensor, source, destination)`
    ///
    /// # Arguments
    /// * `source` - Original positions of dimensions to move
    /// * `destination` - Target positions for the dimensions
    ///
    /// # Examples
    /// ```ignore
    /// let x = Tensor::from_data(vec![1.0; 24], vec![2, 3, 4], DeviceType::Cpu)?;
    /// let y = x.movedim(&[0, 1], &[2, 0])?; // [2,3,4] -> [3,4,2]
    /// ```
    pub fn movedim(&self, source: &[isize], destination: &[isize]) -> Result<Self> {
        if source.len() != destination.len() {
            return Err(TorshError::InvalidArgument(
                "source and destination must have the same length".to_string(),
            ));
        }

        let ndim = self.ndim();

        // Normalize source and destination dimensions
        let norm_source: Result<Vec<usize>> = source
            .iter()
            .map(|&d| {
                let dim = if d < 0 {
                    (ndim as isize + d) as usize
                } else {
                    d as usize
                };
                if dim >= ndim {
                    Err(TorshError::InvalidArgument(format!(
                        "Dimension {} out of range for {}-D tensor",
                        d, ndim
                    )))
                } else {
                    Ok(dim)
                }
            })
            .collect();
        let norm_source = norm_source?;

        let norm_dest: Result<Vec<usize>> = destination
            .iter()
            .map(|&d| {
                let dim = if d < 0 {
                    (ndim as isize + d) as usize
                } else {
                    d as usize
                };
                if dim >= ndim {
                    Err(TorshError::InvalidArgument(format!(
                        "Dimension {} out of range for {}-D tensor",
                        d, ndim
                    )))
                } else {
                    Ok(dim)
                }
            })
            .collect();
        let norm_dest = norm_dest?;

        // Check for duplicates in source
        for i in 0..norm_source.len() {
            for j in i + 1..norm_source.len() {
                if norm_source[i] == norm_source[j] {
                    return Err(TorshError::InvalidArgument(
                        "repeated dim in source".to_string(),
                    ));
                }
            }
        }

        // Check for duplicates in destination
        for i in 0..norm_dest.len() {
            for j in i + 1..norm_dest.len() {
                if norm_dest[i] == norm_dest[j] {
                    return Err(TorshError::InvalidArgument(
                        "repeated dim in destination".to_string(),
                    ));
                }
            }
        }

        // Build permutation array by placing dims in final positions
        let mut result_perm = vec![0; ndim];
        let mut used = vec![false; ndim];

        // Place source dims at destination positions
        for (&src, &dst) in norm_source.iter().zip(norm_dest.iter()) {
            result_perm[dst] = src;
            used[dst] = true;
        }

        // Fill remaining positions with remaining dims in order
        let remaining_dims: Vec<usize> = (0..ndim).filter(|d| !norm_source.contains(d)).collect();

        let mut remaining_idx = 0;
        for i in 0..ndim {
            if !used[i] {
                result_perm[i] = remaining_dims[remaining_idx];
                remaining_idx += 1;
            }
        }

        // Convert usize to i32 for permute
        let perm_i32: Vec<i32> = result_perm.iter().map(|&d| d as i32).collect();
        self.permute(&perm_i32)
    }

    /// Move axis from source position to destination position (alias for movedim)
    ///
    /// # PyTorch Compatibility
    /// Equivalent to `torch.moveaxis(tensor, source, destination)`
    ///
    /// # Arguments
    /// * `source` - Original positions of axes to move
    /// * `destination` - Target positions for the axes
    pub fn moveaxis(&self, source: &[isize], destination: &[isize]) -> Result<Self> {
        self.movedim(source, destination)
    }

    /// Swap two dimensions
    ///
    /// # PyTorch Compatibility
    /// Equivalent to `torch.swapaxes(tensor, axis0, axis1)` or `torch.swapdims(tensor, dim0, dim1)`
    ///
    /// # Arguments
    /// * `axis0` - First dimension
    /// * `axis1` - Second dimension
    ///
    /// # Examples
    /// ```ignore
    /// let x = Tensor::from_data(vec![1.0; 12], vec![2, 3, 2], DeviceType::Cpu)?;
    /// let y = x.swapaxes(0, 2)?; // [2,3,2] -> [2,3,2] with dims 0 and 2 swapped
    /// ```
    pub fn swapaxes(&self, axis0: isize, axis1: isize) -> Result<Self> {
        let ndim = self.ndim();

        // Normalize dimensions
        let dim0 = if axis0 < 0 {
            (ndim as isize + axis0) as usize
        } else {
            axis0 as usize
        };
        let dim1 = if axis1 < 0 {
            (ndim as isize + axis1) as usize
        } else {
            axis1 as usize
        };

        if dim0 >= ndim {
            return Err(TorshError::InvalidArgument(format!(
                "Dimension {} out of range for {}-D tensor",
                axis0, ndim
            )));
        }
        if dim1 >= ndim {
            return Err(TorshError::InvalidArgument(format!(
                "Dimension {} out of range for {}-D tensor",
                axis1, ndim
            )));
        }

        // Build permutation: swap dim0 and dim1
        let mut perm: Vec<i32> = (0..ndim as i32).collect();
        perm.swap(dim0, dim1);

        self.permute(&perm)
    }

    /// Swap two dimensions (alias for swapaxes)
    ///
    /// # PyTorch Compatibility
    /// Equivalent to `torch.swapdims(tensor, dim0, dim1)`
    pub fn swapdims(&self, dim0: isize, dim1: isize) -> Result<Self> {
        self.swapaxes(dim0, dim1)
    }

    /// Broadcast tensor to a new shape
    ///
    /// # PyTorch Compatibility
    /// Equivalent to `torch.broadcast_to(tensor, shape)`
    ///
    /// # Arguments
    /// * `shape` - Target shape for broadcasting
    ///
    /// # Examples
    /// ```ignore
    /// let x = Tensor::from_data(vec![1.0, 2.0], vec![2], DeviceType::Cpu)?;
    /// let y = x.broadcast_to(&[3, 2])?; // Broadcast [2] to [3, 2]
    /// ```
    pub fn broadcast_to(&self, shape: &[usize]) -> Result<Self> {
        // Use the existing expand method which handles broadcasting
        self.expand(shape)
    }

    /// Expand tensor to match another tensor's shape
    ///
    /// # PyTorch Compatibility
    /// Equivalent to `torch.expand_as(tensor, other)`
    ///
    /// # Arguments
    /// * `other` - Target tensor whose shape to match
    ///
    /// # Examples
    /// ```ignore
    /// let x = Tensor::from_data(vec![1.0, 2.0], vec![2], DeviceType::Cpu)?;
    /// let y = Tensor::from_data(vec![0.0; 6], vec![3, 2], DeviceType::Cpu)?;
    /// let z = x.expand_as(&y)?; // Expand x to match y's shape [3, 2]
    /// ```
    pub fn expand_as(&self, other: &Self) -> Result<Self> {
        self.broadcast_to(other.shape().dims())
    }
}

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

    #[test]
    fn test_tensor_view() {
        let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
        let tensor = Tensor::from_data(data, vec![2, 3], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let reshaped = tensor.view(&[3, 2]).expect("view should succeed");
        assert_eq!(reshaped.shape().dims(), &[3, 2]);
        assert_eq!(reshaped.numel(), 6);
    }

    #[test]
    fn test_tensor_view_with_inference() {
        let data = vec![1.0f32; 24];
        let tensor = Tensor::from_data(data, vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let reshaped = tensor.view(&[6, -1]).expect("view should succeed");
        assert_eq!(reshaped.shape().dims(), &[6, 4]);
    }

    #[test]
    fn test_tensor_slice() {
        let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
        let tensor = Tensor::from_data(data, vec![2, 3], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let slice = tensor.slice_tensor(1, 1, 3).expect("slice should succeed");
        assert_eq!(slice.shape().dims(), &[2, 2]);
    }

    #[test]
    fn test_tensor_transpose() {
        let data = vec![1.0f32, 2.0, 3.0, 4.0];
        let tensor = Tensor::from_data(data, vec![2, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let transposed = tensor.transpose(0, 1).expect("transpose should succeed");
        assert_eq!(transposed.shape().dims(), &[2, 2]);
        assert_eq!(
            transposed.get(&[0, 1]).expect("data access should succeed"),
            3.0
        );
        assert_eq!(
            transposed.get(&[1, 0]).expect("data access should succeed"),
            2.0
        );
    }

    #[test]
    fn test_tensor_squeeze_unsqueeze() {
        let data = vec![1.0f32, 2.0, 3.0];
        let tensor = Tensor::from_data(data, vec![1, 3], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let squeezed = tensor.squeeze(0).expect("squeeze should succeed");
        assert_eq!(squeezed.shape().dims(), &[3]);

        let unsqueezed = squeezed.unsqueeze(0).expect("unsqueeze should succeed");
        assert_eq!(unsqueezed.shape().dims(), &[1, 3]);
    }

    #[test]
    fn test_tensor_permute() {
        let data = vec![1.0f32; 24];
        let tensor = Tensor::from_data(data, vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let permuted = tensor.permute(&[2, 0, 1]).expect("permute should succeed");
        assert_eq!(permuted.shape().dims(), &[4, 2, 3]);
    }

    #[test]
    fn test_is_contiguous() {
        let data = vec![1.0f32, 2.0, 3.0, 4.0];
        let tensor = Tensor::from_data(data, vec![2, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");
        assert!(tensor.is_contiguous());

        let transposed = tensor
            .transpose_view(0, 1)
            .expect("transpose view should succeed");
        assert!(!transposed.is_contiguous());

        let contiguous = transposed.contiguous().expect("contiguous should succeed");
        assert!(contiguous.is_contiguous());
    }

    #[test]
    fn test_expand() {
        let data = vec![1.0f32, 2.0];
        let tensor = Tensor::from_data(data, vec![1, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let expanded = tensor.expand(&[3, 2]).expect("expand should succeed");
        assert_eq!(expanded.shape().dims(), &[3, 2]);
        assert_eq!(expanded.numel(), 6);

        // Correctness: broadcasting the size-1 leading dimension repeats the row.
        assert_eq!(
            expanded.to_vec().expect("to_vec should succeed"),
            vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0]
        );

        // The size-1 dimension is broadcast with stride 0; the kept dimension
        // retains its contiguous stride of 1, so the result is a non-contiguous view.
        assert_eq!(expanded.strides(), vec![0, 1]);
        assert!(!expanded.is_contiguous());
        assert!(expanded.is_view());
    }

    #[test]
    fn test_expand_new_leading_dimension() {
        // Expanding [2] -> [3, 2] introduces a brand new leading dimension.
        let tensor = Tensor::from_data(vec![7.0f32, 8.0], vec![2], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let expanded = tensor.expand(&[3, 2]).expect("expand should succeed");
        assert_eq!(expanded.shape().dims(), &[3, 2]);

        // New leading dimension -> stride 0; the original dimension keeps stride 1.
        assert_eq!(expanded.strides(), vec![0, 1]);
        assert_eq!(
            expanded.to_vec().expect("to_vec should succeed"),
            vec![7.0, 8.0, 7.0, 8.0, 7.0, 8.0]
        );

        // Element access flows through the strided view.
        assert_eq!(expanded.get(&[0, 1]).expect("get should succeed"), 8.0);
        assert_eq!(expanded.get(&[2, 0]).expect("get should succeed"), 7.0);
    }

    #[test]
    fn test_expand_zero_copy_no_data_duplication() {
        // A single element expanded to one million logical elements must NOT copy
        // data: the resulting view shares the source's one-element storage.
        let source = Tensor::from_data(vec![5.0f32], vec![1], DeviceType::Cpu)
            .expect("tensor creation should succeed");
        let source_memory = source.memory_usage();

        let expanded = source.expand(&[1024, 1024]).expect("expand should succeed");
        assert_eq!(expanded.numel(), 1024 * 1024);

        // Shared storage => identical reported memory, far below a materialized copy.
        assert_eq!(expanded.memory_usage(), source_memory);
        assert!(expanded.memory_usage() < expanded.numel() * std::mem::size_of::<f32>());

        // Fully broadcast view: both dimensions have stride 0.
        assert_eq!(expanded.strides(), vec![0, 0]);
        assert!(!expanded.is_contiguous());
        assert!(expanded.is_view());

        // Every logical element resolves to the single stored value.
        assert_eq!(expanded.get(&[0, 0]).expect("get should succeed"), 5.0);
        assert_eq!(expanded.get(&[500, 700]).expect("get should succeed"), 5.0);
        assert_eq!(
            expanded.get(&[1023, 1023]).expect("get should succeed"),
            5.0
        );
    }

    #[test]
    fn test_expand_rejects_incompatible_dimension() {
        // Expanding a non-unit dimension to a different size is invalid.
        let tensor = Tensor::from_data(vec![1.0f32, 2.0, 3.0], vec![3], DeviceType::Cpu)
            .expect("tensor creation should succeed");
        assert!(tensor.expand(&[5]).is_err());
        // Fewer dimensions than the source is also invalid.
        let matrix = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0], vec![2, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");
        assert!(matrix.expand(&[4]).is_err());
    }

    #[test]
    fn test_view_error_handling() {
        let data = vec![1.0f32, 2.0, 3.0];
        let tensor = Tensor::from_data(data, vec![3], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        // Should fail - wrong total size
        assert!(tensor.view(&[2, 2]).is_err());

        // Should fail - multiple -1
        assert!(tensor.view(&[-1, -1]).is_err());
    }

    #[test]
    fn test_movedim_single() {
        let tensor = Tensor::from_data(vec![1.0f32; 24], vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        // Move dim 0 to position 2: [2,3,4] -> [3,4,2]
        let result = tensor.movedim(&[0], &[2]).expect("movedim should succeed");
        assert_eq!(result.shape().dims(), &[3, 4, 2]);
    }

    #[test]
    fn test_movedim_multiple() {
        let tensor = Tensor::from_data(vec![1.0f32; 24], vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        // Move dims [0, 1] to positions [2, 0]: [2,3,4] -> [3,4,2]
        let result = tensor
            .movedim(&[0, 1], &[2, 0])
            .expect("movedim should succeed");
        assert_eq!(result.shape().dims(), &[3, 4, 2]);
    }

    #[test]
    fn test_movedim_negative_indices() {
        let tensor = Tensor::from_data(vec![1.0f32; 24], vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        // Move last dim to first position: [2,3,4] -> [4,2,3]
        let result = tensor.movedim(&[-1], &[0]).expect("movedim should succeed");
        assert_eq!(result.shape().dims(), &[4, 2, 3]);
    }

    #[test]
    fn test_moveaxis_alias() {
        let tensor = Tensor::from_data(vec![1.0f32; 24], vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let result1 = tensor.movedim(&[0], &[2]).expect("movedim should succeed");
        let result2 = tensor
            .moveaxis(&[0], &[2])
            .expect("moveaxis should succeed");
        assert_eq!(result1.shape().dims(), result2.shape().dims());
    }

    #[test]
    fn test_swapaxes_simple() {
        let tensor = Tensor::from_data(
            vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0],
            vec![2, 3],
            DeviceType::Cpu,
        )
        .expect("tensor creation should succeed");

        // Swap dims 0 and 1: [2,3] -> [3,2]
        let result = tensor.swapaxes(0, 1).expect("swapaxes should succeed");
        assert_eq!(result.shape().dims(), &[3, 2]);
    }

    #[test]
    fn test_swapaxes_3d() {
        let tensor = Tensor::from_data(vec![1.0f32; 24], vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        // Swap dims 0 and 2: [2,3,4] -> [4,3,2]
        let result = tensor.swapaxes(0, 2).expect("swapaxes should succeed");
        assert_eq!(result.shape().dims(), &[4, 3, 2]);
    }

    #[test]
    fn test_swapaxes_negative_indices() {
        let tensor = Tensor::from_data(vec![1.0f32; 24], vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        // Swap last two dims: [2,3,4] -> [2,4,3]
        let result = tensor.swapaxes(-1, -2).expect("swapaxes should succeed");
        assert_eq!(result.shape().dims(), &[2, 4, 3]);
    }

    #[test]
    fn test_swapdims_alias() {
        let tensor = Tensor::from_data(vec![1.0f32; 24], vec![2, 3, 4], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let result1 = tensor.swapaxes(0, 2).expect("swapaxes should succeed");
        let result2 = tensor.swapdims(0, 2).expect("swapdims should succeed");
        assert_eq!(result1.shape().dims(), result2.shape().dims());
    }

    #[test]
    fn test_broadcast_to_same_shape() {
        let tensor = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0], vec![2, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let result = tensor
            .broadcast_to(&[2, 2])
            .expect("broadcast_to should succeed");
        assert_eq!(result.shape().dims(), &[2, 2]);
    }

    #[test]
    fn test_broadcast_to_expand_dim() {
        let tensor = Tensor::from_data(vec![1.0f32, 2.0], vec![1, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        // Broadcast [1, 2] to [3, 2]
        let result = tensor
            .broadcast_to(&[3, 2])
            .expect("broadcast_to should succeed");
        assert_eq!(result.shape().dims(), &[3, 2]);
    }

    #[test]
    fn test_expand_as_basic() {
        let tensor = Tensor::from_data(vec![1.0f32, 2.0], vec![1, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let target = Tensor::from_data(vec![0.0f32; 6], vec![3, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let result = tensor.expand_as(&target).expect("expand_as should succeed");
        assert_eq!(result.shape().dims(), target.shape().dims());
        assert_eq!(result.shape().dims(), &[3, 2]);
    }
}