etensor_core/

tensor.rs

//! The central Tensor representation and global identity tracking.

use std::sync::atomic::{AtomicU64, Ordering};
use crate::buffer::Buffer;
use crate::device::Device;
use crate::dtypes::DType;
use crate::shape::Shape;

// =====================================================================
// GLOBAL IDENTITY TRACKER
// =====================================================================

/// A static atomic counter ensuring every tensor globally receives a unique ID.
/// This allows the Autograd Tape to track mathematical histories safely across 
/// thousands of simultaneous multi-threaded operations.
static NEXT_TENSOR_ID: AtomicU64 = AtomicU64::new(1);

/// A unique identifier for a Tensor within the computation graph.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TensorId(pub u64);

impl TensorId {
    /// Generates the next globally unique Tensor ID.
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        // Ordering::Relaxed is safe here because we only care about uniqueness,
        // not synchronizing memory access around the ID generation itself.
        TensorId(NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed))
    }
}

// =====================================================================
// THE ATOM
// =====================================================================

/// The core Tensor struct. 
/// 
/// Strictly acts as a metadata wrapper. It points to a physical memory `Buffer` 
/// and maps it using a geometric `Shape`. It contains NO recursive graph pointers.
#[derive(Debug, Clone)]
pub struct Tensor {
    /// The globally unique token used by the Tape to track gradients.
    pub id: TensorId,
    /// An optional string for explainability (e.g., "layer_1_weights").
    pub name: Option<String>,
    /// The physical memory container (wrapped in Arc for zero-copy sharing).
    pub data: Buffer,
    /// The geometric layout and memory strides.
    pub shape: Shape,
    /// The hardware context where the physical memory resides.
    pub device: Device,
    /// The precision format of the memory buffer.
    pub dtype: DType,
    /// Flag indicating whether the Autograd engine should track this tensor.
    pub requires_grad: bool,
}

impl Tensor {
    /// Constructs a new Tensor from raw components.
    /// Automatically assigns a new unique `TensorId`.
    pub fn new(
        data: Buffer,
        shape: Shape,
        device: Device,
        dtype: DType,
        requires_grad: bool,
    ) -> Self {
        Self {
            id: TensorId::new(),
            name: None,
            data,
            shape,
            device,
            dtype,
            requires_grad,
        }
    }

    /// Builder pattern helper to attach an explainability name to the tensor.
    pub fn with_name(mut self, name: &str) -> Self {
        self.name = Some(name.to_string());
        self
    }

    /// Performs an O(1) mathematical transpose.
    /// 
    /// This generates a NEW tensor (with a new ID for the graph) but explicitly 
    /// clones only the `Arc` buffer pointer, never the underlying memory.
    pub fn transpose(&self) -> Self {
        let new_shape = self.shape.transpose();
        
        Self {
            id: TensorId::new(), // New node in the computation graph!
            name: self.name.as_ref().map(|n| format!("{}_T", n)), // Explainability tracker
            data: self.data.clone(), // Zero-copy Arc increment
            shape: new_shape,
            device: self.device,
            dtype: self.dtype,
            requires_grad: self.requires_grad,
        }
    }
}


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

    #[test]
    fn test_tensor_id_uniqueness() {
        let id1 = TensorId::new();
        let id2 = TensorId::new();
        assert_ne!(id1.0, id2.0, "Sequential IDs must not overlap!");
    }

    #[test]
    fn test_atomic_id_thread_safety() {
        // Prove that if 10 threads create tensors at the exact same time, 
        // the atomic counter never drops or duplicates an ID.
        let mut handles = vec![];
        
        for _ in 0..10 {
            handles.push(thread::spawn(|| {
                let mut local_ids = vec![];
                for _ in 0..100 {
                    local_ids.push(TensorId::new().0);
                }
                local_ids
            }));
        }

        let mut all_ids = vec![];
        for handle in handles {
            all_ids.extend(handle.join().unwrap());
        }

        // Sort and check for duplicates
        all_ids.sort_unstable();
        all_ids.dedup();
        
        assert_eq!(all_ids.len(), 1000, "Race condition detected! Duplicate IDs generated.");
    }

    #[test]
    fn test_tensor_explainability_name() {
        let shape = Shape::new(vec![2, 2]);
        let data = Buffer::new_cpu_zeros(4, DType::F32);
        
        let t = Tensor::new(data, shape, Device::Cpu, DType::F32, true)
            .with_name("attention_weights");
            
        assert_eq!(t.name.unwrap(), "attention_weights");
    }

    #[test]
    fn test_zero_copy_transpose_view() {
        let shape = Shape::new(vec![3, 4]);
        let data = Buffer::new_cpu_zeros(12, DType::F32);
        
        let t1 = Tensor::new(data, shape, Device::Cpu, DType::F32, true).with_name("matrix");
        let initial_arc_count = t1.data.strong_count().unwrap();
        
        // Transpose the tensor
        let t2 = t1.transpose();
        
        // 1. Must be a different mathematical node (Different ID)
        assert_ne!(t1.id, t2.id);
        
        // 2. Geometry must be updated
        assert_eq!(t2.shape.dims, vec![4, 3]);
        
        // 3. Explainability tracking must carry over
        assert_eq!(t2.name.unwrap(), "matrix_T");
        
        // 4. Memory MUST NOT move (Arc count increments, physical RAM remains untouched)
        assert_eq!(t2.data.strong_count().unwrap(), initial_arc_count + 1);
    }
}