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, ViewKind};

/// Row-major (C-order) strides for `shape`.
fn contiguous_strides(shape: &[usize]) -> Vec<usize> {
    let mut strides = vec![1usize; shape.len()];
    for axis in (0..shape.len().saturating_sub(1)).rev() {
        strides[axis] = strides[axis + 1] * shape[axis + 1];
    }
    strides
}

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

    /// Strong handle to the tensor a new view of `self` should point at.
    ///
    /// A view of a view keeps pointing at the original base; a base tensor hands
    /// out a handle to itself. The handle is stored by the view, so — unlike the
    /// `Weak` it replaces — it is alive for as long as the view is.
    fn view_base(&self) -> Arc<Self> {
        match &self.base_tensor {
            Some(base) => Arc::clone(base),
            None => Arc::new(self.clone()),
        }
    }

    /// Zero-copy view of `self` with `new_shape`, for a contiguous source.
    ///
    /// The element order is unchanged, so only the metadata differs. `None`
    /// strides mean "plain row-major block starting at the beginning of the
    /// storage" and are therefore only used when the source itself is one; a
    /// contiguous *window* into a larger buffer keeps explicit strides so reads
    /// keep going through the offset-aware path.
    fn contiguous_view_of(&self, new_shape: Vec<usize>) -> Self {
        let strides = if self.strides.is_none() && self.storage_offset == 0 {
            None
        } else {
            Some(contiguous_strides(&new_shape))
        };

        Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: crate::should_record_grad(self.requires_grad),
            grad: Arc::new(RwLock::new(None)), // Views don't share gradients
            operation: Operation::Leaf,        // Overwritten by `record_view`
            strides,
            storage_offset: self.storage_offset,
            base_tensor: Some(self.view_base()),
        }
    }

    /// Record `result` as a shape-only view of `self` in the autograd graph.
    ///
    /// Nothing is recorded when the source does not require gradients, so
    /// inference pipelines keep building plain leaves.
    pub(crate) fn record_view(&self, result: &mut Self, kind: ViewKind) {
        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = Operation::View {
                input: Arc::new(self.clone()),
                kind,
            };
        }
    }

    /// 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.
    ///
    /// A contiguous source is reshaped **without copying**: the result shares its
    /// storage, so writing through either tensor is visible in the other. Only a
    /// non-contiguous source (a transposed or otherwise strided view) is
    /// materialised in view order first. Either way `requires_grad` is
    /// propagated and the reshape is recorded, so gradients keep flowing.
    ///
    /// # 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()
            )));
        }

        // A reshape of a contiguous tensor is pure metadata: the result shares
        // the source's storage (PyTorch `view` semantics). A non-contiguous
        // source has to be materialised in view order first, which is the only
        // case where data is copied.
        let mut result = if self.is_contiguous() {
            self.contiguous_view_of(new_shape)
        } else {
            let data = self.to_vec()?;
            let mut copied = Self::from_data(data, new_shape, self.device)?;
            copied.requires_grad = crate::should_record_grad(self.requires_grad);
            copied
        };
        self.record_view(&mut result, ViewKind::Reshape);
        Ok(result)
    }

    /// 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
        let mut result = self.contiguous_view_of(shape.to_vec());
        self.record_view(&mut result, ViewKind::Reshape);
        Ok(result)
    }

    /// Create a view of a slice along a dimension (shares data, no copying).
    ///
    /// Under `requires_grad` the slice records its geometry
    /// ([`ViewKind::Narrow`]) so gradients scatter back into the source as one
    /// zero-padded slab, which means the source is retained for the graph's
    /// lifetime; inference (no grad) keeps the plain zero-copy view.
    ///
    /// [`Tensor::narrow`] is the same view under PyTorch's argument
    /// convention (signed `dim`/`start`, a `length` instead of an `end`, and an
    /// empty range allowed); the two share one view constructor, so a range
    /// spelled either way yields the same layout and the same gradient.
    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]
            )));
        }

        Ok(self.narrow_view(dim, start, end - start))
    }

    /// Build the aliasing view of a contiguous `length`-long range on `dim`.
    ///
    /// The single construction site for a narrowed view: the storage handle is
    /// shared, the source's strides are kept, and only the storage offset moves
    /// to the start of the window — so the result addresses the *same* buffer
    /// and writes through it (PyTorch's `narrow`/basic-slicing semantics).
    ///
    /// Callers validate `dim`, `start` and `length` under their own error
    /// contract ([`Self::slice_tensor`] and [`Tensor::narrow`] word their
    /// messages differently), so nothing is re-checked here; `start` must be a
    /// valid index on `dim` and `start + length` must not exceed its extent.
    /// A `length` of 0 is a well-formed empty window — it keeps the offset of
    /// its (still in-range) start and simply has no elements to read.
    pub(crate) fn narrow_view(&self, dim: usize, start: usize, length: usize) -> Self {
        let mut new_shape = self.shape.dims().to_vec();
        new_shape[dim] = length;

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

        let mut result = Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: crate::should_record_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: Some(self.view_base()),
        };

        // Record the slice's geometry so gradients scatter back into `self` as
        // one zero-padded slab. The backward pass works in `self`'s *logical*
        // (default row-major) order — not the view's physical strides — because
        // the scatter target is a zeros tensor in `self`'s logical order.
        self.record_view(&mut result, ViewKind::Narrow { dim, start });

        result
    }

    /// 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);

        let mut result = Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: crate::should_record_grad(self.requires_grad),
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: Some(self.view_base()),
        };

        // A transpose is the permutation that swaps the two axes.
        let mut perm: Vec<usize> = (0..self.ndim()).collect();
        perm.swap(dim0, dim1);
        self.record_view(&mut result, ViewKind::Permute(perm));
        Ok(result)
    }

    /// 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);

        let mut result = Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: crate::should_record_grad(self.requires_grad),
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: Some(self.view_base()),
        };
        // Dropping an extent-1 axis keeps the element order, so the gradient is
        // just reshaped back.
        self.record_view(&mut result, ViewKind::Reshape);
        Ok(result)
    }

    /// 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();
        // The inserted axis has extent 1, so any stride addresses the same
        // element — but it must be the *contiguous* one, or a contiguous tensor
        // would start reporting itself as non-contiguous. That is the stride of
        // the axis which now follows it, times that axis' extent (1 when the new
        // axis is appended last).
        let new_stride = if dim == new_shape.len() - 1 {
            1 // Last dimension always has stride 1
        } else {
            new_strides[dim] * new_shape[dim + 1]
        };
        new_strides.insert(dim, new_stride);

        let mut result = Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: crate::should_record_grad(self.requires_grad),
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: Some(self.view_base()),
        };
        // Inserting an extent-1 axis keeps the element order.
        self.record_view(&mut result, ViewKind::Reshape);
        Ok(result)
    }

    /// Transposes two dimensions of the tensor.
    ///
    /// Swaps the specified dimensions. For 2D tensors, calling `transpose(0, 1)`
    /// produces the standard matrix transpose operation.
    ///
    /// The result is a **view** for every rank: it shares storage with the
    /// source and only swaps the two strides, so writing through either tensor
    /// is visible in the other. Call [`Self::contiguous`] when a packed buffer is
    /// needed.
    ///
    /// # 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
            )));
        }

        // Every rank goes through the same path: `transpose` is a view, exactly
        // like PyTorch's. Callers that need a packed buffer call `.contiguous()`.
        self.transpose_view(dim0, dim1)
    }

    /// 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();

        let mut result = Self {
            storage: self.storage.clone(),
            shape: Shape::new(new_shape),
            device: self.device,
            requires_grad: crate::should_record_grad(self.requires_grad),
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: Some(self.view_base()),
        };
        self.record_view(&mut result, ViewKind::Permute(perm_dims));
        Ok(result)
    }

    /// 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
    ///
    /// Dropping extent-1 axes never changes the element order, so this goes
    /// through [`Self::view`]: the result shares storage with a contiguous
    /// source and keeps the tensor connected to the autograd graph. When every
    /// dimension is 1 the result is a scalar (0-dimensional) tensor.
    pub fn squeeze_all(&self) -> Result<Self> {
        let shape = self.shape.dims();
        let new_shape: Result<Vec<i32>> = shape
            .iter()
            .copied()
            .filter(|&s| s != 1)
            .map(|s| {
                i32::try_from(s).map_err(|_| {
                    TorshError::InvalidShape(format!("Dimension {s} is too large to reshape"))
                })
            })
            .collect();

        self.view(&new_shape?)
    }

    /// 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() {
            // Already contiguous: the clone shares the gradient slot, so the graph
            // (and any recorded operation) is preserved unchanged.
            Ok(self.clone())
        } else {
            // The copy reorders nothing logically — output element `i` equals
            // input element `i` in row-major order — so the backward rule is the
            // identity, recorded as a reshape to the (identical) input shape.
            // Without this the copy would reset to `Operation::Leaf` and detach
            // the graph.
            let data = self.to_vec()?;
            let mut result = Self::from_data(data, self.shape.dims().to_vec(), self.device)?;
            self.record_view(&mut result, ViewKind::Reshape);
            Ok(result)
        }
    }

    /// 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.

        let mut result = Self {
            storage: self.storage.clone(),
            shape: Shape::new(shape.to_vec()),
            device: self.device,
            requires_grad: crate::should_record_grad(self.requires_grad),
            grad: Arc::new(RwLock::new(None)),
            operation: Operation::Leaf,
            strides: Some(new_strides),
            storage_offset: self.storage_offset,
            base_tensor: Some(self.view_base()),
        };
        // Backward sums the gradient over the stride-0 (broadcast) axes.
        self.record_view(&mut result, ViewKind::Expand);
        Ok(result)
    }

    /// 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());
    }

    /// F152/F272: a view's base handle must still be usable after the
    /// constructor returns (the old `Arc::downgrade(&Arc::new(..))` idiom made
    /// it dead on arrival), and a view of a view must point at the original.
    #[test]
    fn test_view_constructors_keep_the_base_alive() {
        let base = Tensor::from_data(vec![1.0f32, 2.0], vec![1, 2], DeviceType::Cpu)
            .expect("tensor creation should succeed");

        let expanded = base.expand(&[3, 2]).expect("expand should succeed");
        let recorded = expanded
            .base_tensor
            .as_ref()
            .expect("expand must record its source");
        assert_eq!(recorded.shape().dims(), &[1, 2]);
        assert_eq!(recorded.to_vec().expect("to_vec"), vec![1.0, 2.0]);

        // A view of a view keeps pointing at the original base tensor.
        let chained = expanded
            .transpose_view(0, 1)
            .expect("transpose_view should succeed");
        let chained_base = chained
            .base_tensor
            .as_ref()
            .expect("a chained view must keep a base");
        assert_eq!(chained_base.shape().dims(), &[1, 2]);
        assert!(Arc::ptr_eq(recorded, chained_base));

        // The public accessor exposes the same handle.
        assert!(base.base_tensor().is_none(), "a base tensor has no source");
        assert!(expanded.base_tensor().is_some());
    }

    #[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]);
    }
}