trustformers_core/tensor/
utils.rs

1//! Tensor utility functions.
2//!
3//! This module contains utility functions for working with tensors.
4
5use super::{DType, Tensor};
6use crate::errors::{Result, TrustformersError};
7use scirs2_core::ndarray::{ArrayD, IxDyn};
8use std::collections::HashMap;
9use std::sync::atomic::AtomicU64;
10use std::sync::{Arc, RwLock};
11
12/// Global counter for unique tensor IDs
13#[allow(dead_code)] // Reserved for future tensor tracking features
14static TENSOR_ID_COUNTER: AtomicU64 = AtomicU64::new(0);
15
16lazy_static::lazy_static! {
17    /// Global gradient registry for tracking tensor gradients
18    static ref GRADIENT_REGISTRY: Arc<RwLock<HashMap<u64, Tensor>>> = Arc::new(RwLock::new(HashMap::new()));
19}
20
21thread_local! {
22    /// Thread-local gradient mode flag
23    static GRADIENT_MODE: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
24}
25
26/// Enable gradient tracking for current thread
27pub fn enable_grad() {
28    GRADIENT_MODE.with(|mode| mode.set(true));
29}
30
31/// Disable gradient tracking for current thread
32pub fn disable_grad() {
33    GRADIENT_MODE.with(|mode| mode.set(false));
34}
35
36/// Check if gradient tracking is enabled for current thread
37pub fn is_grad_enabled() -> bool {
38    GRADIENT_MODE.with(|mode| mode.get())
39}
40
41/// Clear all gradients from the registry
42pub fn clear_gradients() {
43    if let Ok(mut registry) = GRADIENT_REGISTRY.write() {
44        registry.clear();
45    }
46}
47
48impl Tensor {
49    /// Get a unique identifier for this tensor instance
50    fn tensor_id(&self) -> u64 {
51        // Generate a hash-based ID from tensor data and shape
52        use std::collections::hash_map::DefaultHasher;
53        use std::hash::{Hash, Hasher};
54
55        let mut hasher = DefaultHasher::new();
56        self.shape().hash(&mut hasher);
57
58        match self {
59            Tensor::F32(arr) => {
60                arr.as_ptr().hash(&mut hasher);
61                arr.len().hash(&mut hasher);
62            },
63            Tensor::F64(arr) => {
64                arr.as_ptr().hash(&mut hasher);
65                arr.len().hash(&mut hasher);
66            },
67            Tensor::I64(arr) => {
68                arr.as_ptr().hash(&mut hasher);
69                arr.len().hash(&mut hasher);
70            },
71            #[cfg(all(target_os = "macos", feature = "metal"))]
72            Tensor::Metal(data) => {
73                // Use buffer_id as a unique identifier for Metal tensors
74                data.buffer_id.hash(&mut hasher);
75                self.len().hash(&mut hasher);
76            },
77            #[cfg(feature = "cuda")]
78            Tensor::CUDA(data) => {
79                // Use the buffer id as a unique identifier for CUDA tensors
80                data.buffer_id().hash(&mut hasher);
81                self.len().hash(&mut hasher);
82            },
83            _ => {
84                // For other tensor types, use a simpler approach
85                self.len().hash(&mut hasher);
86            },
87        }
88
89        hasher.finish()
90    }
91    /// Get the shape of the tensor.
92    ///
93    /// # Returns
94    ///
95    /// A vector containing the dimensions of the tensor.
96    pub fn shape(&self) -> Vec<usize> {
97        match self {
98            Tensor::F32(a) => a.shape().to_vec(),
99            Tensor::F64(a) => a.shape().to_vec(),
100            Tensor::F16(a) => a.shape().to_vec(),
101            Tensor::BF16(a) => a.shape().to_vec(),
102            Tensor::I64(a) => a.shape().to_vec(),
103            Tensor::C32(a) => a.shape().to_vec(),
104            Tensor::C64(a) => a.shape().to_vec(),
105            Tensor::CF16(a) => a.shape().to_vec(),
106            Tensor::CBF16(a) => a.shape().to_vec(),
107            Tensor::Sparse(s) => s.shape().to_vec(),
108            #[cfg(feature = "candle")]
109            Tensor::Candle(t) => t.shape().dims().to_vec(),
110            #[cfg(all(target_os = "macos", feature = "metal"))]
111            Tensor::Metal(data) => data.shape.clone(),
112            #[cfg(feature = "cuda")]
113            Tensor::CUDA(data) => data.shape.clone(),
114        }
115    }
116
117    /// Get the number of elements in the tensor.
118    ///
119    /// # Returns
120    ///
121    /// The total number of elements in the tensor.
122    pub fn len(&self) -> usize {
123        match self {
124            Tensor::F32(a) => a.len(),
125            Tensor::F64(a) => a.len(),
126            Tensor::F16(a) => a.len(),
127            Tensor::BF16(a) => a.len(),
128            Tensor::I64(a) => a.len(),
129            Tensor::C32(a) => a.len(),
130            Tensor::C64(a) => a.len(),
131            Tensor::CF16(a) => a.len(),
132            Tensor::CBF16(a) => a.len(),
133            Tensor::Sparse(s) => s.nnz(), // Non-zero elements for sparse tensors
134            #[cfg(feature = "candle")]
135            Tensor::Candle(t) => t.elem_count(),
136            #[cfg(all(target_os = "macos", feature = "metal"))]
137            Tensor::Metal(data) => data.shape.iter().product(),
138            #[cfg(feature = "cuda")]
139            Tensor::CUDA(data) => data.shape.iter().product(),
140        }
141    }
142
143    /// Check if the tensor is empty.
144    ///
145    /// # Returns
146    ///
147    /// True if the tensor has no elements.
148    pub fn is_empty(&self) -> bool {
149        self.len() == 0
150    }
151
152    /// Get the number of dimensions in the tensor.
153    ///
154    /// # Returns
155    ///
156    /// The number of dimensions.
157    pub fn ndim(&self) -> usize {
158        self.shape().len()
159    }
160
161    /// Get the size in bytes of the tensor.
162    ///
163    /// # Returns
164    ///
165    /// The size in bytes of the tensor data.
166    pub fn size_bytes(&self) -> usize {
167        match self {
168            Tensor::F32(a) => a.len() * std::mem::size_of::<f32>(),
169            Tensor::F64(a) => a.len() * std::mem::size_of::<f64>(),
170            Tensor::F16(a) => a.len() * std::mem::size_of::<half::f16>(),
171            Tensor::BF16(a) => a.len() * std::mem::size_of::<half::bf16>(),
172            Tensor::I64(a) => a.len() * std::mem::size_of::<i64>(),
173            Tensor::C32(a) => a.len() * std::mem::size_of::<scirs2_core::Complex32>(),
174            Tensor::C64(a) => a.len() * std::mem::size_of::<scirs2_core::Complex64>(),
175            Tensor::CF16(a) => a.len() * std::mem::size_of::<scirs2_core::Complex<half::f16>>(),
176            Tensor::CBF16(a) => a.len() * std::mem::size_of::<scirs2_core::Complex<half::bf16>>(),
177            Tensor::Sparse(s) => s.nnz() * std::mem::size_of::<f32>(), // Simplified estimate
178            #[cfg(feature = "candle")]
179            Tensor::Candle(t) => t.elem_count() * std::mem::size_of::<f32>(), // Simplified
180            #[cfg(all(target_os = "macos", feature = "metal"))]
181            Tensor::Metal(data) => {
182                let num_elements: usize = data.shape.iter().product();
183                num_elements * data.dtype.size_in_bytes()
184            },
185            #[cfg(feature = "cuda")]
186            Tensor::CUDA(data) => {
187                let num_elements: usize = data.shape.iter().product();
188                num_elements * data.dtype.size_in_bytes()
189            },
190        }
191    }
192
193    /// Transfer tensor to specified device.
194    ///
195    /// # Arguments
196    ///
197    /// * `device` - Device identifier (e.g., "cpu", "cuda:0", "mps", "tpu:0")
198    ///
199    /// # Returns
200    ///
201    /// A tensor on the specified device (currently CPU-only with validation).
202    pub fn to_device(&self, device: &str) -> Result<Tensor> {
203        // Validate device string format and provide helpful error messages
204        let device_lower = device.to_lowercase();
205
206        // Parse device components
207        let (device_type, device_index) = if device_lower.contains(':') {
208            let parts: Vec<&str> = device_lower.split(':').collect();
209            if parts.len() != 2 {
210                return Err(TrustformersError::tensor_op_error(
211                    &format!("Invalid device format '{}'. Expected format: 'device_type' or 'device_type:index'", device),
212                    "to_device"
213                ));
214            }
215
216            let index = parts[1].parse::<usize>().map_err(|_| {
217                TrustformersError::tensor_op_error(
218                    &format!(
219                        "Invalid device index '{}'. Expected a non-negative integer",
220                        parts[1]
221                    ),
222                    "to_device",
223                )
224            })?;
225
226            (parts[0], Some(index))
227        } else {
228            (device_lower.as_str(), None)
229        };
230
231        // Validate supported device types
232        match device_type {
233            "cpu" => {
234                // CPU is always supported
235                if let Some(index) = device_index {
236                    if index > 0 {
237                        return Err(TrustformersError::tensor_op_error(
238                            &format!("CPU device index {} not supported. CPU only supports index 0 or no index", index),
239                            "to_device"
240                        ));
241                    }
242                }
243                // Return a clone for CPU (no actual transfer needed)
244                Ok(self.clone())
245            },
246            "cuda" => {
247                // CUDA support would require additional backend integration
248                if let Some(index) = device_index {
249                    Err(TrustformersError::tensor_op_error(
250                        &format!("CUDA device cuda:{} not available. This build doesn't support CUDA. Consider using CPU instead with device='cpu'", index),
251                        "to_device"
252                    ))
253                } else {
254                    Err(TrustformersError::tensor_op_error(
255                        "CUDA devices not available. This build doesn't support CUDA. Consider using CPU instead with device='cpu'",
256                        "to_device"
257                    ))
258                }
259            },
260            "mps" => {
261                // Metal Performance Shaders (Apple Silicon)
262                Err(TrustformersError::tensor_op_error(
263                    "MPS device not available. This build doesn't support Metal Performance Shaders. Consider using CPU instead with device='cpu'",
264                    "to_device"
265                ))
266            },
267            "tpu" => {
268                // Tensor Processing Unit
269                Err(TrustformersError::tensor_op_error(
270                    "TPU devices not available. This build doesn't support TPU. Consider using CPU instead with device='cpu'",
271                    "to_device"
272                ))
273            },
274            "xpu" | "intel" => {
275                // Intel XPU (Intel GPU/AI accelerators)
276                Err(TrustformersError::tensor_op_error(
277                    "Intel XPU devices not available. This build doesn't support Intel XPU. Consider using CPU instead with device='cpu'",
278                    "to_device"
279                ))
280            },
281            "npu" => {
282                // Neural Processing Unit
283                Err(TrustformersError::tensor_op_error(
284                    "NPU devices not available. This build doesn't support NPU. Consider using CPU instead with device='cpu'",
285                    "to_device"
286                ))
287            },
288            _ => {
289                Err(TrustformersError::tensor_op_error(
290                    &format!("Unknown device type '{}'. Supported device types: cpu, cuda, mps, tpu, xpu, npu. For this build, only 'cpu' is supported", device_type),
291                    "to_device"
292                ))
293            },
294        }
295    }
296
297    /// Transfer tensor to specified device using Device enum.
298    ///
299    /// This is the preferred method for device transfers in modern code.
300    /// It supports Metal GPU acceleration and provides better type safety.
301    ///
302    /// # Arguments
303    ///
304    /// * `device` - Device enum (Device::CPU, Device::Metal(0), etc.)
305    ///
306    /// # Returns
307    ///
308    /// A tensor on the specified device.
309    ///
310    /// # Example
311    ///
312    /// ```no_run
313    /// use trustformers_core::tensor::Tensor;
314    /// use trustformers_core::device::Device;
315    ///
316    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
317    /// let cpu_tensor = Tensor::randn(&[2, 3])?;
318    ///
319    /// // Transfer to Metal GPU
320    /// let gpu_tensor = cpu_tensor.to_device_enum(&Device::Metal(0))?;
321    ///
322    /// // Transfer back to CPU
323    /// let result = gpu_tensor.to_device_enum(&Device::CPU)?;
324    /// # Ok(())
325    /// # }
326    /// ```
327    pub fn to_device_enum(&self, device: &crate::device::Device) -> Result<Tensor> {
328        match (self, device) {
329            // F32 → Metal
330            #[cfg(all(target_os = "macos", feature = "metal"))]
331            (Tensor::F32(arr), crate::device::Device::Metal(_)) => {
332                use crate::gpu_ops::metal::get_metal_backend;
333                let backend = get_metal_backend()?;
334                let data_vec: Vec<f32> = arr.iter().copied().collect();
335
336                #[cfg(debug_assertions)]
337                {
338                    // Debug: Verify data_vec before GPU upload (only in debug builds)
339                    eprintln!(
340                        "🔍 to_device_enum(F32→Metal): data_vec.len()={}",
341                        data_vec.len()
342                    );
343                    if !data_vec.is_empty() {
344                        eprintln!(
345                            "🔍 to_device_enum: first 10 values: {:?}",
346                            &data_vec[..10.min(data_vec.len())]
347                        );
348                        eprintln!(
349                            "🔍 to_device_enum: stats - min={:.4}, max={:.4}, mean={:.4}",
350                            data_vec.iter().fold(f32::INFINITY, |a, &b| a.min(b)),
351                            data_vec.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)),
352                            data_vec.iter().sum::<f32>() / data_vec.len() as f32
353                        );
354                    }
355                }
356
357                let buffer_id = backend.create_persistent_buffer(&data_vec)?;
358
359                #[cfg(debug_assertions)]
360                {
361                    eprintln!("🔍 to_device_enum: Created buffer_id={:?}", buffer_id);
362
363                    // Verify by immediately downloading (GPU→CPU transfer - expensive!)
364                    let verify_data = backend.download_buffer_to_vec(&buffer_id)?;
365                    eprintln!(
366                        "🔍 to_device_enum: Verification download - len={}, first 10: {:?}",
367                        verify_data.len(),
368                        &verify_data[..10.min(verify_data.len())]
369                    );
370                }
371
372                Ok(Tensor::Metal(super::MetalTensorData {
373                    buffer_id,
374                    shape: arr.shape().to_vec(),
375                    dtype: DType::F32,
376                }))
377            },
378
379            // F64 → Metal (convert to F32 first)
380            #[cfg(all(target_os = "macos", feature = "metal"))]
381            (Tensor::F64(arr), crate::device::Device::Metal(_)) => {
382                use crate::gpu_ops::metal::get_metal_backend;
383                let backend = get_metal_backend()?;
384                let data_vec: Vec<f32> = arr.iter().map(|&x| x as f32).collect();
385                let buffer_id = backend.create_persistent_buffer(&data_vec)?;
386                Ok(Tensor::Metal(super::MetalTensorData {
387                    buffer_id,
388                    shape: arr.shape().to_vec(),
389                    dtype: DType::F32,
390                }))
391            },
392
393            // Metal → F32
394            #[cfg(all(target_os = "macos", feature = "metal"))]
395            (Tensor::Metal(metal_data), crate::device::Device::CPU) => {
396                use crate::gpu_ops::metal::get_metal_backend;
397                let backend = get_metal_backend()?;
398                let buffer = backend.get_persistent_buffer(&metal_data.buffer_id)?;
399
400                // Download from GPU
401                let size: usize = metal_data.shape.iter().product();
402
403                // Handle different dtypes
404                match metal_data.dtype {
405                    DType::F32 => {
406                        let ptr = buffer.contents() as *const f32;
407                        let data_vec = unsafe { std::slice::from_raw_parts(ptr, size) }.to_vec();
408
409                        // Convert to ArrayD
410                        use scirs2_core::ndarray::ArrayD;
411                        let arr = ArrayD::from_shape_vec(
412                            scirs2_core::ndarray::IxDyn(&metal_data.shape),
413                            data_vec,
414                        )
415                        .map_err(|e| {
416                            TrustformersError::tensor_op_error(
417                                &format!("Failed to create array from shape: {}", e),
418                                "to_device_enum",
419                            )
420                        })?;
421                        Ok(Tensor::F32(arr))
422                    },
423                    _ => Err(TrustformersError::tensor_op_error(
424                        &format!("Unsupported Metal tensor dtype: {:?}", metal_data.dtype),
425                        "to_device_enum",
426                    )),
427                }
428            },
429
430            // Metal → Metal (different device, currently just clone)
431            #[cfg(all(target_os = "macos", feature = "metal"))]
432            (Tensor::Metal(metal_data), crate::device::Device::Metal(_)) => {
433                // For now, just return a clone (buffer is reference counted)
434                // TODO: Implement actual device-to-device transfer if needed
435                Ok(Tensor::Metal(metal_data.clone()))
436            },
437
438            // Already on correct device - no-op
439            (Tensor::F32(_), crate::device::Device::CPU) => Ok(self.clone()),
440            (Tensor::F64(_), crate::device::Device::CPU) => Ok(self.clone()),
441            (Tensor::F16(_), crate::device::Device::CPU) => Ok(self.clone()),
442            (Tensor::BF16(_), crate::device::Device::CPU) => Ok(self.clone()),
443            (Tensor::I64(_), crate::device::Device::CPU) => Ok(self.clone()),
444            (Tensor::C32(_), crate::device::Device::CPU) => Ok(self.clone()),
445            (Tensor::C64(_), crate::device::Device::CPU) => Ok(self.clone()),
446            (Tensor::CF16(_), crate::device::Device::CPU) => Ok(self.clone()),
447            (Tensor::CBF16(_), crate::device::Device::CPU) => Ok(self.clone()),
448            (Tensor::Sparse(_), crate::device::Device::CPU) => Ok(self.clone()),
449
450            // Metal not available in this build
451            #[cfg(not(feature = "metal"))]
452            (_, crate::device::Device::Metal(_)) => Err(TrustformersError::hardware_error(
453                "Metal not available. Compile with --features metal",
454                "to_device_enum",
455            )),
456
457            // F32 → CUDA
458            #[cfg(feature = "cuda")]
459            #[allow(unused_variables)]
460            (Tensor::F32(arr), crate::device::Device::CUDA(device_id)) => {
461                #[cfg(any(target_os = "linux", target_os = "windows"))]
462                {
463                    use crate::gpu_ops::cuda::get_cuda_backend;
464                    let backend = get_cuda_backend(*device_id)?;
465                    let data_vec: Vec<f32> = arr.iter().copied().collect();
466                    let buffer_id = backend.create_persistent_buffer(&data_vec)?;
467                    Ok(Tensor::CUDA(super::CudaTensorData::new(
468                        buffer_id,
469                        *device_id,
470                        arr.shape().to_vec(),
471                        DType::F32,
472                    )))
473                }
474                #[cfg(not(any(target_os = "linux", target_os = "windows")))]
475                {
476                    Err(TrustformersError::hardware_error(
477                        "CUDA is only supported on Linux and Windows",
478                        "to_device_enum",
479                    ))
480                }
481            },
482
483            // F64 → CUDA (convert to F32 first)
484            #[cfg(feature = "cuda")]
485            #[allow(unused_variables)]
486            (Tensor::F64(arr), crate::device::Device::CUDA(device_id)) => {
487                #[cfg(any(target_os = "linux", target_os = "windows"))]
488                {
489                    use crate::gpu_ops::cuda::get_cuda_backend;
490                    let backend = get_cuda_backend(*device_id)?;
491                    let data_vec: Vec<f32> = arr.iter().map(|&x| x as f32).collect();
492                    let buffer_id = backend.create_persistent_buffer(&data_vec)?;
493                    Ok(Tensor::CUDA(super::CudaTensorData::new(
494                        buffer_id,
495                        *device_id,
496                        arr.shape().to_vec(),
497                        DType::F32,
498                    )))
499                }
500                #[cfg(not(any(target_os = "linux", target_os = "windows")))]
501                {
502                    Err(TrustformersError::hardware_error(
503                        "CUDA is only supported on Linux and Windows",
504                        "to_device_enum",
505                    ))
506                }
507            },
508
509            // CUDA → F32
510            #[cfg(feature = "cuda")]
511            #[allow(unused_variables)]
512            (Tensor::CUDA(cuda_data), crate::device::Device::CPU) => {
513                #[cfg(any(target_os = "linux", target_os = "windows"))]
514                {
515                    use crate::gpu_ops::cuda::get_cuda_backend;
516                    // Download from the device the buffer actually lives on (the
517                    // handle carries the ordinal), not a hardcoded device 0.
518                    let backend = get_cuda_backend(cuda_data.device_id())?;
519
520                    // Handle different dtypes
521                    match cuda_data.dtype {
522                        DType::F32 => {
523                            // Download from GPU to CPU
524                            let data_vec = backend.download_buffer(&cuda_data.buffer_id())?;
525
526                            // Convert to ArrayD
527                            use scirs2_core::ndarray::ArrayD;
528                            let arr = ArrayD::from_shape_vec(
529                                scirs2_core::ndarray::IxDyn(&cuda_data.shape),
530                                data_vec,
531                            )
532                            .map_err(|e| {
533                                TrustformersError::tensor_op_error(
534                                    &format!("Failed to create array from shape: {}", e),
535                                    "to_device_enum",
536                                )
537                            })?;
538                            Ok(Tensor::F32(arr))
539                        },
540                        _ => Err(TrustformersError::tensor_op_error(
541                            &format!("Unsupported CUDA tensor dtype: {:?}", cuda_data.dtype),
542                            "to_device_enum",
543                        )),
544                    }
545                }
546                #[cfg(not(any(target_os = "linux", target_os = "windows")))]
547                {
548                    Err(TrustformersError::hardware_error(
549                        "CUDA is only supported on Linux and Windows",
550                        "to_device_enum",
551                    ))
552                }
553            },
554
555            // CUDA → CUDA
556            #[cfg(feature = "cuda")]
557            (Tensor::CUDA(cuda_data), crate::device::Device::CUDA(target_device)) => {
558                if *target_device == cuda_data.device_id() {
559                    // Same device: the clone shares the same reference-counted
560                    // resident buffer (no copy, refcount increment only).
561                    Ok(Tensor::CUDA(cuda_data.clone()))
562                } else {
563                    // Cross-device transfer: bounce through the host (download from
564                    // the source device, re-upload to the target device).
565                    let host_tensor = self.to_device_enum(&crate::device::Device::CPU)?;
566                    host_tensor.to_device_enum(device)
567                }
568            },
569
570            // CUDA not available in this build
571            #[cfg(not(feature = "cuda"))]
572            (_, crate::device::Device::CUDA(_)) => Err(TrustformersError::hardware_error(
573                "CUDA not available. Compile with --features cuda",
574                "to_device_enum",
575            )),
576
577            // ROCm transfers (placeholder for future implementation)
578            (_, crate::device::Device::ROCm(_)) => Err(TrustformersError::hardware_error(
579                "ROCm transfer not implemented yet",
580                "to_device_enum",
581            )),
582
583            // WebGPU transfers (placeholder for future implementation)
584            (_, crate::device::Device::WebGPU) => Err(TrustformersError::hardware_error(
585                "WebGPU transfer not implemented yet",
586                "to_device_enum",
587            )),
588
589            // Fallback for unsupported combinations
590            #[allow(unreachable_patterns)]
591            _ => Err(TrustformersError::tensor_op_error(
592                &format!(
593                    "Unsupported device transfer from {:?} to {:?}",
594                    self.dtype(),
595                    device
596                ),
597                "to_device_enum",
598            )),
599        }
600    }
601
602    /// Get gradient tensor.
603    ///
604    /// # Returns
605    ///
606    /// Returns the gradient tensor associated with this tensor, or None if no gradient exists.
607    /// Gradients are only tracked when gradient mode is enabled via `enable_grad()`.
608    ///
609    /// # Example
610    ///
611    /// ```no_run
612    /// use trustformers_core::tensor::{Tensor, enable_grad, disable_grad};
613    ///
614    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
615    /// enable_grad();
616    /// let x = Tensor::randn(&[2, 3])?;
617    /// // After some computation that requires gradients...
618    /// if let Ok(grad_tensor) = x.grad() {
619    ///     println!("Gradient: {:?}", grad_tensor.shape());
620    /// }
621    /// disable_grad();
622    /// # Ok(())
623    /// # }
624    /// ```
625    pub fn grad(&self) -> Result<Tensor> {
626        if !is_grad_enabled() {
627            return Err(TrustformersError::tensor_op_error(
628                "Gradient tracking is not enabled. Use enable_grad() to enable gradient tracking.",
629                "grad",
630            ));
631        }
632
633        let tensor_id = self.tensor_id();
634
635        if let Ok(registry) = GRADIENT_REGISTRY.read() {
636            if let Some(grad_tensor) = registry.get(&tensor_id) {
637                Ok(grad_tensor.clone())
638            } else {
639                Err(TrustformersError::tensor_op_error(
640                    "No gradient found for this tensor. Gradients are set during backward pass.",
641                    "grad",
642                ))
643            }
644        } else {
645            Err(TrustformersError::tensor_op_error(
646                "Failed to access gradient registry.",
647                "grad",
648            ))
649        }
650    }
651
652    /// Set gradient tensor.
653    ///
654    /// # Arguments
655    ///
656    /// * `grad` - The gradient tensor to set for this tensor
657    ///
658    /// # Returns
659    ///
660    /// Returns Ok(()) if the gradient was successfully set, or an error if gradient tracking
661    /// is not enabled or if the gradient shape doesn't match the tensor shape.
662    ///
663    /// # Example
664    ///
665    /// ```no_run
666    /// use trustformers_core::tensor::{Tensor, enable_grad};
667    ///
668    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
669    /// enable_grad();
670    /// let mut x = Tensor::randn(&[2, 3])?;
671    /// let grad = Tensor::ones(&[2, 3])?;
672    /// x.set_grad(grad)?;
673    /// # Ok(())
674    /// # }
675    /// ```
676    pub fn set_grad(&mut self, grad: Tensor) -> Result<()> {
677        if !is_grad_enabled() {
678            return Err(TrustformersError::tensor_op_error(
679                "Gradient tracking is not enabled. Use enable_grad() to enable gradient tracking.",
680                "set_grad",
681            ));
682        }
683
684        // Validate gradient shape matches tensor shape
685        if self.shape() != grad.shape() {
686            return Err(TrustformersError::tensor_op_error(
687                &format!(
688                    "Gradient shape {:?} doesn't match tensor shape {:?}",
689                    grad.shape(),
690                    self.shape()
691                ),
692                "set_grad",
693            ));
694        }
695
696        let tensor_id = self.tensor_id();
697
698        if let Ok(mut registry) = GRADIENT_REGISTRY.write() {
699            registry.insert(tensor_id, grad);
700            Ok(())
701        } else {
702            Err(TrustformersError::tensor_op_error(
703                "Failed to access gradient registry.",
704                "set_grad",
705            ))
706        }
707    }
708
709    /// Get tensor data as a vector (for F32 tensors).
710    ///
711    /// # Returns
712    ///
713    /// A Result containing a vector with the tensor data.
714    pub fn data(&self) -> Result<Vec<f32>> {
715        match self {
716            Tensor::F32(a) => Ok(a.iter().cloned().collect()),
717            Tensor::F64(a) => Ok(a.iter().map(|&x| x as f32).collect()),
718            Tensor::I64(a) => Ok(a.iter().map(|&x| x as f32).collect()),
719            #[cfg(all(target_os = "macos", feature = "metal"))]
720            Tensor::Metal(_) => {
721                // Convert to CPU first, then get data
722                let cpu_tensor = self.to_device_enum(&crate::device::Device::CPU)?;
723                cpu_tensor.data()
724            },
725            #[cfg(feature = "cuda")]
726            Tensor::CUDA(_) => {
727                // Convert to CPU first, then get data
728                let cpu_tensor = self.to_device_enum(&crate::device::Device::CPU)?;
729                cpu_tensor.data()
730            },
731            _ => Err(TrustformersError::tensor_op_error(
732                "Unsupported tensor type for data conversion",
733                "data_conversion",
734            )),
735        }
736    }
737
738    /// Normalised Shannon entropy of the softmax distribution over the tensor values.
739    ///
740    /// Values are mapped to a probability distribution with a numerically-stable softmax,
741    /// then `H = -∑ p·ln(p)` is divided by `ln(n)` so the result lands in `[0, 1]`
742    /// (≈0 when one value dominates, ≈1 when the values are uniform). Useful as an
743    /// uncertainty signal for adaptive / early-exit computation.
744    pub fn softmax_entropy_normalized(&self) -> Result<f32> {
745        let values = self.to_vec_f32()?;
746        let n = values.len();
747        if n <= 1 {
748            return Ok(0.0);
749        }
750        let max = values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
751        let exps: Vec<f32> = values.iter().map(|&v| (v - max).exp()).collect();
752        let sum: f32 = exps.iter().sum();
753        if sum <= 0.0 {
754            return Ok(0.0);
755        }
756        let entropy: f32 =
757            exps.iter().map(|&e| e / sum).filter(|&p| p > 1e-10).map(|p| -p * p.ln()).sum();
758        Ok((entropy / (n as f32).ln()).clamp(0.0, 1.0))
759    }
760
761    /// Get tensor data as F32 vector (alias for data() method).
762    ///
763    /// # Returns
764    ///
765    /// A Result containing a vector with the tensor data as f32.
766    pub fn data_f32(&self) -> Result<Vec<f32>> {
767        self.data()
768    }
769
770    /// Set tensor data from F32 vector.
771    ///
772    /// # Arguments
773    ///
774    /// * `data` - Vector of f32 values to set as tensor data
775    ///
776    /// # Returns
777    ///
778    /// A Result indicating success or failure.
779    pub fn set_data_f32(&mut self, data: &[f32]) -> Result<()> {
780        match self {
781            Tensor::F32(a) => {
782                let shape = a.shape().to_vec();
783                let expected_len: usize = shape.iter().product();
784                if data.len() != expected_len {
785                    return Err(TrustformersError::tensor_op_error(
786                        &format!(
787                            "Data length {} does not match tensor size {}",
788                            data.len(),
789                            expected_len
790                        ),
791                        "set_data_f32",
792                    ));
793                }
794                *a = ArrayD::from_shape_vec(IxDyn(&shape), data.to_vec()).map_err(|e| {
795                    TrustformersError::tensor_op_error(&e.to_string(), "set_data_f32")
796                })?;
797                Ok(())
798            },
799            _ => Err(TrustformersError::tensor_op_error(
800                "set_data_f32 only supported for F32 tensors",
801                "set_data_f32",
802            )),
803        }
804    }
805
806    /// Get mutable reference to tensor data as a slice (for F32 tensors).
807    ///
808    /// # Returns
809    ///
810    /// A Result containing a mutable slice of the tensor data.
811    pub fn data_mut(&mut self) -> Result<&mut [f32]> {
812        match self {
813            Tensor::F32(a) => a.as_slice_mut().ok_or_else(|| {
814                TrustformersError::tensor_op_error(
815                    "Tensor data must be contiguous for mutable access",
816                    "data_mut",
817                )
818            }),
819            _ => Err(TrustformersError::tensor_op_error(
820                "Mutable data access only supported for F32 tensors",
821                "data_mut",
822            )),
823        }
824    }
825
826    /// Modify tensor data in-place with a closure.
827    ///
828    /// # Arguments
829    ///
830    /// * `f` - A closure that takes a mutable slice of the tensor data
831    ///
832    /// # Returns
833    ///
834    /// A Result indicating success or failure.
835    pub fn modify_data<F>(&mut self, f: F) -> Result<()>
836    where
837        F: FnOnce(&mut [f32]),
838    {
839        match self {
840            Tensor::F32(a) => {
841                if let Some(slice) = a.as_slice_mut() {
842                    f(slice);
843                    Ok(())
844                } else {
845                    Err(TrustformersError::tensor_op_error(
846                        "Cannot get mutable slice",
847                        "modify_data",
848                    ))
849                }
850            },
851            _ => Err(TrustformersError::tensor_op_error(
852                "Modify data only supported for F32 tensors",
853                "modify_data",
854            )),
855        }
856    }
857
858    /// Get the device where the tensor is stored.
859    ///
860    /// # Returns
861    ///
862    /// A string representing the device.
863    pub fn device(&self) -> String {
864        match self {
865            Tensor::F32(_)
866            | Tensor::F64(_)
867            | Tensor::F16(_)
868            | Tensor::BF16(_)
869            | Tensor::I64(_)
870            | Tensor::C32(_)
871            | Tensor::C64(_)
872            | Tensor::CF16(_)
873            | Tensor::CBF16(_) => "cpu".to_string(),
874            Tensor::Sparse(_) => "cpu".to_string(),
875            #[cfg(feature = "candle")]
876            Tensor::Candle(t) => format!("{:?}", t.device()),
877            #[cfg(all(target_os = "macos", feature = "metal"))]
878            Tensor::Metal(_) => "metal".to_string(),
879            #[cfg(feature = "cuda")]
880            Tensor::CUDA(_) => "cuda".to_string(),
881        }
882    }
883
884    /// Get the number of elements in the tensor.
885    ///
886    /// # Returns
887    ///
888    /// The total number of elements.
889    pub fn size(&self) -> usize {
890        self.shape().iter().product()
891    }
892
893    /// Get memory usage in bytes.
894    ///
895    /// # Returns
896    ///
897    /// Memory usage in bytes.
898    pub fn memory_usage(&self) -> usize {
899        match self {
900            Tensor::F32(a) => a.len() * std::mem::size_of::<f32>(),
901            Tensor::F64(a) => a.len() * std::mem::size_of::<f64>(),
902            Tensor::F16(a) => a.len() * std::mem::size_of::<half::f16>(),
903            Tensor::BF16(a) => a.len() * std::mem::size_of::<half::bf16>(),
904            Tensor::I64(a) => a.len() * std::mem::size_of::<i64>(),
905            Tensor::C32(a) => a.len() * std::mem::size_of::<scirs2_core::Complex32>(),
906            Tensor::C64(a) => a.len() * std::mem::size_of::<scirs2_core::Complex64>(),
907            Tensor::CF16(a) => a.len() * std::mem::size_of::<scirs2_core::Complex<half::f16>>(),
908            Tensor::CBF16(a) => a.len() * std::mem::size_of::<scirs2_core::Complex<half::bf16>>(),
909            Tensor::Sparse(s) => s.memory_usage(),
910            #[cfg(feature = "candle")]
911            Tensor::Candle(t) => t.elem_count() * 4, // Approximate
912            #[cfg(all(target_os = "macos", feature = "metal"))]
913            Tensor::Metal(m) => m.shape.iter().product::<usize>() * 4, // Approximate as f32
914            #[cfg(feature = "cuda")]
915            Tensor::CUDA(c) => c.shape.iter().product::<usize>() * 4, // Approximate as f32
916        }
917    }
918
919    /// Get the data type of the tensor.
920    ///
921    /// # Returns
922    ///
923    /// The data type.
924    pub fn dtype(&self) -> DType {
925        match self {
926            Tensor::F32(_) => DType::F32,
927            Tensor::F64(_) => DType::F64,
928            Tensor::F16(_) => DType::F16,
929            Tensor::BF16(_) => DType::BF16,
930            Tensor::I64(_) => DType::I64,
931            Tensor::C32(_) => DType::C32,
932            Tensor::C64(_) => DType::C64,
933            Tensor::CF16(_) => DType::CF16,
934            Tensor::CBF16(_) => DType::CBF16,
935            Tensor::Sparse(_) => DType::F32, // Sparse tensors use f32 by default
936            #[cfg(feature = "candle")]
937            Tensor::Candle(_) => DType::F32, // Default assumption
938            #[cfg(all(target_os = "macos", feature = "metal"))]
939            Tensor::Metal(data) => data.dtype,
940            #[cfg(feature = "cuda")]
941            Tensor::CUDA(data) => data.dtype,
942        }
943    }
944
945    /// Get the data type (alias for dtype).
946    pub fn get_dtype(&self) -> DType {
947        self.dtype()
948    }
949
950    /// Get a float value at a specific index.
951    ///
952    /// # Arguments
953    ///
954    /// * `index` - The linear index
955    ///
956    /// # Returns
957    ///
958    /// The float value at the index.
959    pub fn get_float(&self, index: usize) -> Result<f32> {
960        match self {
961            Tensor::F32(a) => {
962                if index >= a.len() {
963                    return Err(TrustformersError::tensor_op_error(
964                        &format!(
965                            "Index {} out of bounds for tensor of size {}",
966                            index,
967                            a.len()
968                        ),
969                        "get_float",
970                    ));
971                }
972                Ok(a.iter().nth(index).copied().unwrap_or(0.0))
973            },
974            Tensor::F64(a) => {
975                if index >= a.len() {
976                    return Err(TrustformersError::tensor_op_error(
977                        &format!(
978                            "Index {} out of bounds for tensor of size {}",
979                            index,
980                            a.len()
981                        ),
982                        "get_float",
983                    ));
984                }
985                Ok(a.iter().nth(index).copied().unwrap_or(0.0) as f32)
986            },
987            #[cfg(all(target_os = "macos", feature = "metal"))]
988            Tensor::Metal(_) => {
989                // Convert to CPU first, then get float
990                let cpu_tensor = self.to_device_enum(&crate::device::Device::CPU)?;
991                cpu_tensor.get_float(index)
992            },
993            #[cfg(feature = "cuda")]
994            Tensor::CUDA(_) => {
995                // Convert to CPU first, then get float
996                let cpu_tensor = self.to_device_enum(&crate::device::Device::CPU)?;
997                cpu_tensor.get_float(index)
998            },
999            _ => Err(TrustformersError::tensor_op_error(
1000                "Get float not supported for this tensor type",
1001                "get_float",
1002            )),
1003        }
1004    }
1005
1006    /// Get a scalar value from a 0-dimensional or 1-element tensor.
1007    ///
1008    /// # Type Parameters
1009    ///
1010    /// * `T` - The type to convert to (i32, i64, f32, f64)
1011    ///
1012    /// # Returns
1013    ///
1014    /// The scalar value.
1015    pub fn item<T>(&self) -> Result<T>
1016    where
1017        T: num_traits::NumCast,
1018    {
1019        if self.len() != 1 {
1020            return Err(TrustformersError::tensor_op_error(
1021                &format!(
1022                    "item() requires a single-element tensor, but got {} elements",
1023                    self.len()
1024                ),
1025                "item",
1026            ));
1027        }
1028
1029        match self {
1030            Tensor::F32(a) => {
1031                let val = a.iter().next().copied().unwrap_or(0.0);
1032                T::from(val).ok_or_else(|| {
1033                    TrustformersError::tensor_op_error(
1034                        "Failed to convert f32 to target type",
1035                        "item",
1036                    )
1037                })
1038            },
1039            Tensor::F64(a) => {
1040                let val = a.iter().next().copied().unwrap_or(0.0);
1041                T::from(val).ok_or_else(|| {
1042                    TrustformersError::tensor_op_error(
1043                        "Failed to convert f64 to target type",
1044                        "item",
1045                    )
1046                })
1047            },
1048            Tensor::I64(a) => {
1049                let val = a.iter().next().copied().unwrap_or(0);
1050                T::from(val).ok_or_else(|| {
1051                    TrustformersError::tensor_op_error(
1052                        "Failed to convert i64 to target type",
1053                        "item",
1054                    )
1055                })
1056            },
1057            #[cfg(all(target_os = "macos", feature = "metal"))]
1058            Tensor::Metal(_) => {
1059                // Convert to CPU first, then get item
1060                let cpu_tensor = self.to_device_enum(&crate::device::Device::CPU)?;
1061                cpu_tensor.item::<T>()
1062            },
1063            #[cfg(feature = "cuda")]
1064            Tensor::CUDA(_) => {
1065                // Convert to CPU first, then get item
1066                let cpu_tensor = self.to_device_enum(&crate::device::Device::CPU)?;
1067                cpu_tensor.item::<T>()
1068            },
1069            _ => Err(TrustformersError::tensor_op_error(
1070                "item() not supported for this tensor type",
1071                "item",
1072            )),
1073        }
1074    }
1075
1076    /// Get an i64 scalar value from a tensor.
1077    ///
1078    /// # Returns
1079    ///
1080    /// The i64 scalar value.
1081    pub fn get_scalar_i64(&self) -> Result<i64> {
1082        self.item::<i64>()
1083    }
1084
1085    /// Compare tensor elements with a scalar value.
1086    /// TEMPORARY: Uses ndarray. Will be replaced with SciRS2-Core.
1087    ///
1088    /// # Arguments
1089    ///
1090    /// * `scalar` - The scalar value to compare against
1091    ///
1092    /// # Returns
1093    ///
1094    /// A boolean tensor where True indicates elements equal to the scalar.
1095    pub fn eq_scalar(&self, scalar: f64) -> Result<Tensor> {
1096        match self {
1097            Tensor::F32(a) => {
1098                let scalar_f32 = scalar as f32;
1099                let result =
1100                    a.mapv(|x| if (x - scalar_f32).abs() < 1e-6 { 1.0f32 } else { 0.0f32 });
1101                Ok(Tensor::F32(result))
1102            },
1103            Tensor::F64(a) => {
1104                let result = a.mapv(|x| if (x - scalar).abs() < 1e-9 { 1.0f64 } else { 0.0f64 });
1105                Ok(Tensor::F64(result))
1106            },
1107            Tensor::I64(a) => {
1108                let scalar_i64 = scalar as i64;
1109                let result = a.mapv(|x| if x == scalar_i64 { 1i64 } else { 0i64 });
1110                Ok(Tensor::I64(result))
1111            },
1112            _ => Err(TrustformersError::tensor_op_error(
1113                "eq_scalar not supported for this tensor type",
1114                "eq_scalar",
1115            )),
1116        }
1117    }
1118
1119    /// Split tensor into batches along the first dimension
1120    ///
1121    /// # Arguments
1122    ///
1123    /// * `batch_size` - Size of each batch
1124    ///
1125    /// # Returns
1126    ///
1127    /// Vector of tensors, each representing a batch. The last batch may be smaller
1128    /// if the tensor size is not evenly divisible by batch_size.
1129    ///
1130    /// # Example
1131    ///
1132    /// ```
1133    /// use trustformers_core::tensor::Tensor;
1134    ///
1135    /// let tensor = Tensor::ones(&[10, 4]).expect("Failed to create ones tensor");
1136    /// let batches = tensor.batch_split(3).unwrap();
1137    /// assert_eq!(batches.len(), 4); // [3, 3, 3, 1]
1138    /// assert_eq!(batches[0].shape(), &[3, 4]);
1139    /// assert_eq!(batches[3].shape(), &[1, 4]);
1140    /// ```
1141    pub fn batch_split(&self, batch_size: usize) -> Result<Vec<Tensor>> {
1142        if batch_size == 0 {
1143            return Err(TrustformersError::tensor_op_error(
1144                "Batch size must be greater than 0",
1145                "batch_split",
1146            ));
1147        }
1148
1149        let shape = self.shape();
1150        if shape.is_empty() {
1151            return Err(TrustformersError::tensor_op_error(
1152                "Cannot batch split a scalar tensor",
1153                "batch_split",
1154            ));
1155        }
1156
1157        let total_size = shape[0];
1158        let mut batches = Vec::new();
1159
1160        for start in (0..total_size).step_by(batch_size) {
1161            let end = std::cmp::min(start + batch_size, total_size);
1162            let batch = self.slice(0, start, end)?;
1163            batches.push(batch);
1164        }
1165
1166        Ok(batches)
1167    }
1168
1169    /// Batch tensors together along a new first dimension
1170    ///
1171    /// # Arguments
1172    ///
1173    /// * `tensors` - Slice of tensors to batch together. All tensors must have the same shape.
1174    ///
1175    /// # Returns
1176    ///
1177    /// A new tensor with shape [batch_size, original_shape...]
1178    ///
1179    /// # Example
1180    ///
1181    /// ```
1182    /// use trustformers_core::tensor::Tensor;
1183    ///
1184    /// let t1 = Tensor::ones(&[3, 4]).expect("Failed to create ones tensor");
1185    /// let t2 = Tensor::zeros(&[3, 4]).expect("Failed to create zero tensor");
1186    /// let t3 = Tensor::ones(&[3, 4]).expect("Failed to create ones tensor");
1187    ///
1188    /// let batched = Tensor::batch_stack(&[&t1, &t2, &t3]).unwrap();
1189    /// assert_eq!(batched.shape(), &[3, 3, 4]);
1190    /// ```
1191    pub fn batch_stack(tensors: &[&Tensor]) -> Result<Tensor> {
1192        if tensors.is_empty() {
1193            return Err(TrustformersError::tensor_op_error(
1194                "Cannot stack empty tensor list",
1195                "batch_stack",
1196            ));
1197        }
1198
1199        // Verify all tensors have the same shape
1200        let reference_shape = tensors[0].shape();
1201        for (i, tensor) in tensors.iter().enumerate() {
1202            if tensor.shape() != reference_shape {
1203                return Err(TrustformersError::tensor_op_error(
1204                    &format!(
1205                        "Tensor {} has shape {:?}, expected {:?}",
1206                        i,
1207                        tensor.shape(),
1208                        reference_shape
1209                    ),
1210                    "batch_stack",
1211                ));
1212            }
1213        }
1214
1215        // Create new shape with batch dimension
1216        let mut new_shape = vec![tensors.len()];
1217        new_shape.extend_from_slice(&reference_shape);
1218
1219        match tensors[0] {
1220            Tensor::F32(_) => {
1221                let mut result_data = Vec::new();
1222                for tensor in tensors {
1223                    if let Tensor::F32(arr) = tensor {
1224                        result_data.extend(arr.iter().copied());
1225                    }
1226                }
1227                Tensor::from_vec(result_data, &new_shape)
1228            },
1229            _ => Err(TrustformersError::tensor_op_error(
1230                "Batch stacking currently only implemented for F32 tensors",
1231                "batch_stack",
1232            )),
1233        }
1234    }
1235
1236    /// Unbatch a tensor by removing the first dimension
1237    ///
1238    /// # Returns
1239    ///
1240    /// Vector of tensors, each representing an item from the batch
1241    ///
1242    /// # Example
1243    ///
1244    /// ```
1245    /// use trustformers_core::tensor::Tensor;
1246    ///
1247    /// let batched = Tensor::ones(&[3, 4, 5]).expect("Failed to create ones tensor");
1248    /// let unbatched = batched.unbatch().unwrap();
1249    /// assert_eq!(unbatched.len(), 3);
1250    /// assert_eq!(unbatched[0].shape(), &[4, 5]);
1251    /// ```
1252    pub fn unbatch(&self) -> Result<Vec<Tensor>> {
1253        let shape = self.shape();
1254        if shape.is_empty() {
1255            return Err(TrustformersError::tensor_op_error(
1256                "Cannot unbatch a scalar tensor",
1257                "unbatch",
1258            ));
1259        }
1260
1261        let batch_size = shape[0];
1262        let mut items = Vec::with_capacity(batch_size);
1263
1264        for i in 0..batch_size {
1265            let item = self.slice(0, i, i + 1)?;
1266            // Remove the batch dimension by squeezing the first axis
1267            let squeezed = item.squeeze(0)?;
1268            items.push(squeezed);
1269        }
1270
1271        Ok(items)
1272    }
1273}
1274
1275#[cfg(test)]
1276mod tests {
1277    use super::*;
1278
1279    #[test]
1280    fn test_gradient_tracking_basic() {
1281        // Test basic gradient functionality
1282        enable_grad();
1283
1284        let mut x = Tensor::ones(&[2, 3]).expect("Failed to create ones tensor");
1285        let grad = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3])
1286            .expect("Tensor from_vec failed");
1287
1288        // Set gradient
1289        assert!(x.set_grad(grad.clone()).is_ok());
1290
1291        // Get gradient
1292        let retrieved_grad = x.grad().expect("operation failed in test");
1293        assert_eq!(retrieved_grad.shape(), vec![2, 3]);
1294
1295        disable_grad();
1296    }
1297
1298    #[test]
1299    fn test_gradient_tracking_disabled() {
1300        // Test that gradients fail when tracking is disabled
1301        disable_grad();
1302
1303        let mut x = Tensor::ones(&[2, 3]).expect("Failed to create ones tensor");
1304        let grad = Tensor::ones(&[2, 3]).expect("Failed to create ones tensor");
1305
1306        // Should fail when gradient tracking is disabled
1307        assert!(x.set_grad(grad).is_err());
1308        assert!(x.grad().is_err());
1309    }
1310
1311    #[test]
1312    fn test_gradient_shape_validation() {
1313        enable_grad();
1314
1315        let mut x = Tensor::ones(&[2, 3]).expect("Failed to create ones tensor");
1316        let wrong_shape_grad = Tensor::ones(&[3, 2]).expect("Failed to create ones tensor");
1317
1318        // Should fail when gradient shape doesn't match tensor shape
1319        assert!(x.set_grad(wrong_shape_grad).is_err());
1320
1321        disable_grad();
1322    }
1323
1324    #[test]
1325    fn test_clear_gradients() {
1326        enable_grad();
1327
1328        let mut x = Tensor::ones(&[2, 3]).expect("Failed to create ones tensor");
1329        let grad = Tensor::ones(&[2, 3]).expect("Failed to create ones tensor");
1330
1331        // Set gradient
1332        x.set_grad(grad).expect("operation failed in test");
1333
1334        // Verify gradient exists
1335        assert!(x.grad().is_ok());
1336
1337        // Clear all gradients
1338        clear_gradients();
1339
1340        // Gradient should no longer exist
1341        assert!(x.grad().is_err());
1342
1343        disable_grad();
1344    }
1345
1346    #[test]
1347    fn test_gradient_mode_functions() {
1348        // Test gradient mode control functions
1349        disable_grad();
1350        assert!(!is_grad_enabled());
1351
1352        enable_grad();
1353        assert!(is_grad_enabled());
1354
1355        disable_grad();
1356        assert!(!is_grad_enabled());
1357    }
1358
1359    #[test]
1360    fn test_softmax_entropy_normalized_bounds() {
1361        // Uniform values → maximal (≈1.0) normalised entropy.
1362        let uniform =
1363            Tensor::from_vec(vec![1.0, 1.0, 1.0, 1.0], &[4]).expect("failed to build tensor");
1364        let h_uniform = uniform.softmax_entropy_normalized().expect("entropy failed");
1365        assert!(
1366            h_uniform > 0.99,
1367            "uniform entropy should be ~1.0, got {h_uniform}"
1368        );
1369
1370        // One value dominates → near-zero normalised entropy.
1371        let peaked =
1372            Tensor::from_vec(vec![20.0, 0.0, 0.0, 0.0], &[4]).expect("failed to build tensor");
1373        let h_peaked = peaked.softmax_entropy_normalized().expect("entropy failed");
1374        assert!(
1375            (0.0..0.2).contains(&h_peaked),
1376            "peaked entropy should be small, got {h_peaked}"
1377        );
1378
1379        // Single element → defined as 0.
1380        let single = Tensor::from_vec(vec![5.0], &[1]).expect("failed to build tensor");
1381        assert_eq!(
1382            single.softmax_entropy_normalized().expect("entropy failed"),
1383            0.0
1384        );
1385    }
1386}
trustformers_core/tensor/utils.rs

trustformers_core/tensor/
utils.rs
