train_station/tensor/indexing/

select.rs

use crate::gradtrack::{is_grad_enabled, GradEngine, GradFn};
use crate::tensor::core::Tensor;

impl Tensor {
    /// Select a slice along a given dimension at a specific index
    ///
    /// This operation extracts a slice from the input tensor by fixing a specific dimension
    /// at a given index. The result is a tensor with one fewer dimension than the input,
    /// containing the selected slice.
    ///
    /// The select operation returns a view (zero-copy) when the base offset is zero,
    /// otherwise it creates a contiguous copy to ensure correctness. This operation is
    /// commonly used for extracting specific rows, columns, or slices from tensors.
    ///
    /// # Arguments
    ///
    /// * `dim` - The dimension along which to select (must be < tensor rank)
    /// * `index` - The index within the specified dimension to select (must be < dim size)
    ///
    /// # Returns
    ///
    /// A tensor with the selected slice. The result has the same shape as the input
    /// except with the specified dimension removed.
    ///
    /// # Performance
    ///
    /// Returns a view when possible (base offset is zero) to avoid copying. On
    /// non-zero offsets, falls back to a contiguous copy for correctness. Gradients
    /// propagate back to the selected slice when GradTrack is enabled.
    ///
    /// # Examples
    ///
    /// ## Basic Row Selection
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Create a 2x3 tensor: [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
    /// let tensor = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![2, 3]).unwrap();
    ///
    /// // Select row 1 (dimension 0, index 1)
    /// let result = tensor.select(0, 1);
    ///
    /// // Result shape is [3] (dimension 0 removed)
    /// assert_eq!(result.shape().dims(), vec![3]);
    /// assert_eq!(result.get(&[0]), 3.0);  // First element of row 1
    /// assert_eq!(result.get(&[1]), 4.0);  // Second element of row 1
    /// assert_eq!(result.get(&[2]), 5.0);  // Third element of row 1
    /// ```
    ///
    /// ## Column Selection
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Create a 2x3 tensor: [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
    /// let tensor = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![2, 3]).unwrap();
    ///
    /// // Select column 1 (dimension 1, index 1)
    /// let result = tensor.select(1, 1);
    ///
    /// // Result shape is [2] (dimension 1 removed)
    /// assert_eq!(result.shape().dims(), vec![2]);
    /// assert_eq!(result.get(&[0]), 1.0);  // Column 1, row 0
    /// assert_eq!(result.get(&[1]), 4.0);  // Column 1, row 1
    /// ```
    ///
    /// ## Select with Gradient Tracking
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// let tensor = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0], vec![2, 2]).unwrap()
    ///     .with_requires_grad();
    ///
    /// // Select row 1 with gradient tracking enabled
    /// let mut result = tensor.select(0, 1);
    /// result.backward(None);
    ///
    /// // Verify gradients are computed correctly
    /// let grad = tensor.grad_owned().expect("gradient missing");
    /// assert_eq!(grad.shape().dims(), vec![2, 2]);
    /// // Only row 1 receives gradients
    /// assert_eq!(grad.get(&[0, 0]), 0.0);  // Row 0: no gradient
    /// assert_eq!(grad.get(&[0, 1]), 0.0);  // Row 0: no gradient
    /// assert_eq!(grad.get(&[1, 0]), 1.0);  // Row 1: gradient flows
    /// assert_eq!(grad.get(&[1, 1]), 1.0);  // Row 1: gradient flows
    /// ```
    ///
    /// # Performance Characteristics
    ///
    /// - **Time Complexity**: O(n) where n is the number of elements in the selected slice
    /// - **Memory Usage**: Zero-copy view when base offset is zero, otherwise creates a copy
    /// - **Optimization**: Uses efficient stride-based access for non-contiguous tensors
    /// - **GradTrack Overhead**: Minimal overhead when gradient tracking is enabled
    /// - **Memory Layout**: Result is contiguous when a copy is made, view otherwise
    ///
    /// # Implementation Details
    ///
    /// The select operation works by:
    /// 1. Validating the dimension and index bounds
    /// 2. Computing the new shape by removing the selected dimension
    /// 3. Computing the new strides by removing the selected dimension's stride
    /// 4. Calculating the base offset for the selected slice
    /// 5. If base offset is zero: creating a view with adjusted shape/strides
    /// 6. If base offset is non-zero: creating a contiguous copy of the slice
    /// 7. Registering the operation for gradient computation if needed
    ///
    /// # Safety
    ///
    /// This function performs comprehensive bounds checking to ensure:
    /// - The tensor has non-zero rank
    /// - The specified dimension is within the tensor's rank
    /// - The index is within bounds for the specified dimension
    /// - Memory access is safe through proper offset calculations
    ///
    /// # Panics
    ///
    /// This function will panic if:
    /// - The tensor has zero rank
    /// - `dim` is greater than or equal to the tensor's rank
    /// - `index` is greater than or equal to the size of the specified dimension
    ///
    /// # Thread Safety
    ///
    /// This function is thread-safe and can be called concurrently on different tensors.
    /// The operation does not modify the input tensor and creates either a view or a new tensor.
    ///
    /// # View vs Copy Behavior
    ///
    /// - **View (zero-copy)**: When the base offset is zero, returns a view that shares
    ///   the same memory as the input tensor with adjusted shape and strides
    /// - **Copy**: When the base offset is non-zero, creates a contiguous copy to ensure
    ///   correctness across all operations
    ///
    /// # GradTrack Behavior
    ///
    /// When gradient tracking is enabled:
    /// - Gradients are scattered back to the selected slice in the input tensor
    /// - Other positions in the input tensor receive zero gradients
    /// - This behavior ensures correct gradient flow for the selected elements
    #[track_caller]
    pub fn select(&self, dim: usize, index: usize) -> Tensor {
        let rank = self.shape().rank();
        assert!(rank > 0, "select requires non-zero rank");
        assert!(
            dim < rank,
            "select dim {} out of bounds for rank {}",
            dim,
            rank
        );
        let dim_size = self.shape().dims()[dim];
        assert!(
            index < dim_size,
            "select index {} out of bounds for dimension {} (size {})",
            index,
            dim,
            dim_size
        );

        // Build new dims/strides removing the selected dimension
        let mut new_dims = Vec::with_capacity(rank - 1);
        let mut new_strides = Vec::with_capacity(rank - 1);
        for i in 0..rank {
            if i == dim {
                continue;
            }
            new_dims.push(self.shape().dims()[i]);
            new_strides.push(self.strides()[i]);
        }

        // Base pointer shift by index * stride(dim)
        let base_offset = index * self.stride(dim);

        // Create zero-copy view using core as_strided with storage_offset
        let mut result = match crate::tensor::core::view::as_strided_view(
            self,
            &new_dims,
            &new_strides,
            base_offset,
        ) {
            Ok(v) => v,
            Err(e) => panic!("select view error: {:?}", e),
        };

        // GradTrack registration only when gradients are enabled and input requires grad
        if self.requires_grad() && is_grad_enabled() {
            result.set_requires_grad(true);
            let grad_fn = GradFn::Select {
                dim,
                index,
                input_shape: self.shape().dims().to_vec(),
            };
            result.set_grad_fn(grad_fn.clone());
            GradEngine::register_operation(result.id(), vec![self.id()], grad_fn);
        }

        result
    }
}

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

    #[test]
    fn test_select_basic() {
        let x = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![2, 3]).unwrap();
        let s = x.select(0, 1);
        assert_eq!(s.shape().dims(), vec![3]);
        assert_eq!(s.get(&[0]), 3.0);
        assert_eq!(s.get(&[2]), 5.0);
    }

    #[test]
    fn test_select_grad() {
        let x = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0], vec![2, 2])
            .unwrap()
            .with_requires_grad();
        let mut s = x.select(0, 1);
        s.backward(None);
        let gx = x.grad_owned().expect("grad missing");
        // Only row 1 receives ones
        assert_eq!(gx.get(&[0, 0]), 0.0);
        assert_eq!(gx.get(&[0, 1]), 0.0);
        assert_eq!(gx.get(&[1, 0]), 1.0);
        assert_eq!(gx.get(&[1, 1]), 1.0);
    }
}