train_station/tensor/transform/

reshape.rs

//! Tensor reshape operations
//!
//! This module provides tensor reshape functionality that changes the
//! dimensions of a tensor while preserving the total number of elements.
//! Reshaping is a fundamental tensor transformation operation used in
//! machine learning for preparing data for different layer types,
//! implementing complex tensor manipulations, and adapting tensor shapes
//! for specific operations.
//!
//! # Operations
//!
//! * `reshape()` - Reshape tensor to specified dimensions with automatic inference
//!
//! # Performance Characteristics
//!
//! * **Zero-Copy Operation**: Returns a view when tensor is contiguous
//! * **Memory Efficient**: Reuses existing tensor data through view operations
//! * **Automatic Inference**: Supports -1 dimension for automatic size calculation
//! * **Gradient Tracking**: Full GradTrack support for automatic differentiation
//! * **Validation**: Comprehensive error checking for invalid reshape operations
//!
//! # Examples
//!
//! ```
//! use train_station::Tensor;
//!
//! // Basic reshape
//! let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
//! let reshaped = tensor.reshape(vec![3, 2]);
//! assert_eq!(reshaped.shape().dims(), vec![3, 2]);
//!
//! // Automatic dimension inference with -1
//! let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
//! let reshaped = tensor.reshape(vec![2, -1]);
//! assert_eq!(reshaped.shape().dims(), vec![2, 2]);
//! ```

use crate::gradtrack::{GradEngine, GradFn};
use crate::tensor::core::Tensor;
// Shape is referenced via core view helpers; direct import not needed here

impl Tensor {
    /// Reshape the tensor to the specified dimensions
    ///
    /// Changes the shape of the tensor while preserving the total number of elements.
    /// This operation returns a view when the tensor is contiguous, avoiding data
    /// copying. For non-contiguous tensors, data is copied to ensure the reshape
    /// is valid.
    ///
    /// The reshape operation supports automatic dimension inference using -1,
    /// which allows one dimension to be automatically calculated based on the
    /// total number of elements and the other specified dimensions.
    ///
    /// # Arguments
    ///
    /// * `new_shape` - Target shape for the tensor. Use -1 for one dimension
    ///   to have it automatically inferred from the total size.
    ///
    /// # Returns
    ///
    /// A new tensor with the specified shape containing the same data as the
    /// original tensor.
    ///
    /// # Panics
    ///
    /// * If more than one dimension is -1
    /// * If the total number of elements doesn't match the original tensor
    /// * If any dimension size is 0 or less than -1
    /// * If the inferred dimension size is not a whole number
    ///
    /// # Examples
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Basic reshape
    /// let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
    /// let reshaped = tensor.reshape(vec![3, 2]);
    /// assert_eq!(reshaped.shape().dims(), vec![3, 2]);
    /// assert_eq!(reshaped.get(&[0, 0]), 1.0);
    /// assert_eq!(reshaped.get(&[2, 1]), 6.0);
    /// ```
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Using -1 for automatic dimension inference
    /// let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
    /// let reshaped = tensor.reshape(vec![2, -1]);
    /// assert_eq!(reshaped.shape().dims(), vec![2, 2]);
    /// assert_eq!(reshaped.get(&[0, 0]), 1.0);
    /// assert_eq!(reshaped.get(&[1, 1]), 4.0);
    /// ```
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Reshape with gradient tracking
    /// let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    /// tensor.set_requires_grad(true);
    ///
    /// let reshaped = tensor.reshape(vec![4]);
    /// assert!(reshaped.requires_grad());
    /// assert_eq!(reshaped.shape().dims(), vec![4]);
    /// ```
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Reshape 3D tensor
    /// let data: Vec<f32> = (0..24).map(|i| i as f32).collect();
    /// let tensor = Tensor::from_slice(&data, vec![2, 3, 4]).unwrap();
    /// let reshaped = tensor.reshape(vec![6, 4]);
    /// assert_eq!(reshaped.shape().dims(), vec![6, 4]);
    /// assert_eq!(reshaped.size(), 24);
    /// ```
    ///
    /// # Performance
    ///
    /// - **Contiguous tensors**: O(1) time complexity, returns a view
    /// - **Non-contiguous tensors**: O(n) time complexity with data copying
    /// - **Memory usage**: No additional allocation for view operations
    /// - **Gradient tracking**: Preserves gradient requirements and tracking
    ///
    /// # Automatic Dimension Inference
    ///
    /// When using -1 for a dimension, the size is automatically calculated:
    /// ```rust
    /// use train_station::Tensor;
    ///
    /// // For a tensor with 12 elements
    /// let data: Vec<f32> = (0..12).map(|i| i as f32).collect();
    /// let tensor = Tensor::from_slice(&data, vec![3, 4]).unwrap();
    ///
    /// let reshaped1 = tensor.reshape(vec![3, -1]);  // Results in shape [3, 4]
    /// let reshaped2 = tensor.reshape(vec![-1, 6]);  // Results in shape [2, 6]
    /// let reshaped3 = tensor.reshape(vec![-1]);     // Results in shape [12]
    /// ```
    #[track_caller]
    pub fn reshape(&self, new_shape: Vec<i32>) -> Tensor {
        // Validate and process the new shape
        let processed_shape = self.process_reshape_dimensions(new_shape);

        // Validate that total size matches
        let new_size: usize = processed_shape.iter().product();
        assert_eq!(
            new_size,
            self.size(),
            "Cannot reshape tensor of size {} to shape {:?} (size {})",
            self.size(),
            processed_shape,
            new_size
        );
        // Zero-copy reshape using core reshape_view validation
        match crate::tensor::core::view::reshape_view(self, &processed_shape) {
            Ok(mut v) => {
                if self.requires_grad() {
                    v.set_requires_grad(true);
                    let grad_fn = GradFn::Reshape {
                        original_shape: self.shape().dims().to_vec(),
                    };
                    v.set_grad_fn(grad_fn.clone());
                    GradEngine::register_operation(v.id(), vec![self.id()], grad_fn);
                }
                v
            }
            Err(_) => {
                // If core reshape fails (non-contiguous), materialize contiguous then view
                let contiguous = self.contiguous();
                let mut v =
                    match crate::tensor::core::view::reshape_view(&contiguous, &processed_shape) {
                        Ok(v) => v,
                        Err(e) => panic!("reshape error: {:?}", e),
                    };
                if self.requires_grad() {
                    v.set_requires_grad(true);
                    let grad_fn = GradFn::Reshape {
                        original_shape: self.shape().dims().to_vec(),
                    };
                    v.set_grad_fn(grad_fn.clone());
                    GradEngine::register_operation(v.id(), vec![self.id()], grad_fn);
                }
                v
            }
        }
    }

    /// Process reshape dimensions and handle -1 inference
    ///
    /// Validates reshape dimensions and automatically infers the size of any
    /// dimension marked as -1. This method ensures that the reshape operation
    /// is valid and calculates the appropriate dimension sizes.
    ///
    /// # Arguments
    ///
    /// * `new_shape` - Target shape with possible -1 for inference
    ///
    /// # Returns
    ///
    /// Processed shape with all dimensions as positive usize values
    ///
    /// # Panics
    ///
    /// * If more than one dimension is -1
    /// * If any dimension size is 0 or less than -1
    /// * If the total size is not divisible by the known dimensions
    ///
    /// # Examples
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
    /// // This internally calls process_reshape_dimensions
    /// let reshaped = tensor.reshape(vec![2, -1]);
    /// assert_eq!(reshaped.shape().dims(), vec![2, 2]);
    /// ```
    pub(crate) fn process_reshape_dimensions(&self, new_shape: Vec<i32>) -> Vec<usize> {
        // Validate input dimensions
        let mut infer_dim = None;
        let mut known_size = 1usize;

        for (i, &dim) in new_shape.iter().enumerate() {
            if dim == -1 {
                if infer_dim.is_some() {
                    panic!("Only one dimension can be -1 for automatic inference");
                }
                infer_dim = Some(i);
            } else if dim <= 0 {
                panic!("Dimension sizes must be positive, got {}", dim);
            } else {
                known_size *= dim as usize;
            }
        }

        // Convert to usize and infer -1 dimension
        let mut processed: Vec<usize> = new_shape
            .iter()
            .map(|&d| if d == -1 { 0 } else { d as usize })
            .collect();

        if let Some(infer_idx) = infer_dim {
            let total_size = self.size();
            if known_size == 0 || !total_size.is_multiple_of(known_size) {
                panic!(
                    "Cannot infer dimension size: total size {} not divisible by known size {}",
                    total_size, known_size
                );
            }
            processed[infer_idx] = total_size / known_size;
        }

        processed
    }

    // removed: private reshape_view in favor of core view::reshape_view
}

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

    #[test]
    fn test_basic_reshape() {
        let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
        let reshaped = tensor.reshape(vec![3, 2]);

        assert_eq!(reshaped.shape().dims(), vec![3, 2]);
        assert_eq!(reshaped.size(), 6);

        // Verify data integrity
        assert_eq!(reshaped.get(&[0, 0]), 1.0);
        assert_eq!(reshaped.get(&[2, 1]), 6.0);
    }

    #[test]
    fn test_reshape_with_inference() {
        let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
        let reshaped = tensor.reshape(vec![2, -1]);

        assert_eq!(reshaped.shape().dims(), vec![2, 2]);
        assert_eq!(reshaped.size(), 4);
    }

    #[test]
    fn test_reshape_autograd() {
        let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
        tensor.set_requires_grad(true);

        let reshaped = tensor.reshape(vec![4]);
        assert!(reshaped.requires_grad());
        assert!(!matches!(reshaped.grad_fn(), GradFn::None));
    }

    #[test]
    #[should_panic(expected = "Only one dimension can be -1")]
    fn test_multiple_infer_dimensions() {
        let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
        tensor.reshape(vec![-1, -1]);
    }

    #[test]
    #[should_panic(expected = "Cannot reshape tensor of size 4")]
    fn test_invalid_size() {
        let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
        tensor.reshape(vec![2, 3]); // 2*3 = 6 != 4
    }

    #[test]
    #[should_panic(expected = "Dimension sizes must be positive")]
    fn test_negative_dimension() {
        let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
        tensor.reshape(vec![2, -2]);
    }

    #[test]
    fn test_large_tensor_reshape() {
        // Test with larger tensor to verify performance
        let data: Vec<f32> = (0..1000).map(|i| i as f32).collect();
        let tensor = Tensor::from_slice(&data, vec![10, 100]).unwrap();

        let reshaped = tensor.reshape(vec![25, 40]);
        assert_eq!(reshaped.shape().dims(), vec![25, 40]);
        assert_eq!(reshaped.size(), 1000);

        // Verify first and last elements preserved
        assert_eq!(reshaped.get(&[0, 0]), 0.0);
        assert_eq!(reshaped.get(&[24, 39]), 999.0);
    }

    #[test]
    fn test_reshape_edge_cases() {
        // Scalar to 1D
        let scalar = Tensor::from_slice(&[42.0], vec![1]).unwrap();
        let reshaped = scalar.reshape(vec![-1]);
        assert_eq!(reshaped.shape().dims(), vec![1]);

        // 1D to scalar (well, size-1 tensor)
        let tensor = Tensor::from_slice(&[42.0], vec![1]).unwrap();
        let reshaped = tensor.reshape(vec![1]);
        assert_eq!(reshaped.shape().dims(), vec![1]);
    }

    #[test]
    fn test_multi_operation_with_reshape() {
        // Test that reshape works well with other operations
        let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
        let reshaped = tensor.reshape(vec![4]);
        let transposed = reshaped.reshape(vec![1, 4]);

        assert_eq!(transposed.shape().dims(), vec![1, 4]);
        assert_eq!(transposed.get(&[0, 3]), 4.0);
    }

    #[test]
    fn test_reshape_with_autograd_chain() {
        // Test autograd with reshape in a computation chain
        let mut a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
        let b = Tensor::from_slice(&[0.5, 0.5, 0.5, 0.5], vec![4]).unwrap();

        a.set_requires_grad(true);

        // Reshape a to match b's shape, then add
        let reshaped_a = a.reshape(vec![4]);
        assert!(reshaped_a.requires_grad());

        let result = reshaped_a.add_tensor_optimized(&b);

        assert_eq!(result.shape().dims(), vec![4]);
        // Note: add_tensor_optimized may not preserve gradients for mixed operations
        // In a full implementation, we'd use the AutogradTensor trait methods

        // Verify values
        assert_eq!(result.get(&[0]), 1.5); // 1.0 + 0.5
        assert_eq!(result.get(&[3]), 4.5); // 4.0 + 0.5
    }
}