torsh_autograd/

lib.rs

//! Automatic differentiation engine for ToRSh
//!
//! This crate provides a PyTorch-compatible autograd API that fully leverages
//! scirs2-autograd's powerful automatic differentiation capabilities.
//!
//! # Quick Start
//!
//! ```rust,ignore
//! use torsh_autograd::prelude::*;
//!
//! // Enable gradient computation
//! let x = tensor::ones(&[2, 3]).requires_grad_(true);
//! let y = x.pow(2).sum();
//!
//! // Compute gradients
//! y.backward();
//! let grad = x.grad();
//! ```
//!
//! # Architecture
//!
//! The autograd system is built around several key components:
//!
//! - **Gradient computation**: Automatic computation of gradients through computation graphs
//! - **Tensor operations**: Differentiable tensor operations with gradient tracking
//! - **Variable management**: Thread-local variable environments for gradient storage
//! - **Guard system**: RAII guards for gradient mode management
//! - **Anomaly detection**: Detection and recovery from numerical anomalies
//! - **SciRS2 integration**: Deep integration with the SciRS2 autograd system
//! - **Hardware acceleration**: Multi-platform support (CUDA, Metal, WebGPU)
//!
//! # API Stability
//!
//! The crate follows semantic versioning with clearly defined stability levels:
//!
//! - **Stable APIs** ([`stable_api::stable`]): Core functionality with backward compatibility
//! - **Beta APIs** ([`stable_api::beta`]): Feature-complete but may evolve
//! - **Experimental APIs** ([`stable_api::experimental`]): May change significantly
//!
//! See [`stable_api`] for details on stability guarantees.
//!
//! # Examples
//!
//! The [`examples`] module provides comprehensive usage examples:
//!
//! - Basic gradient computation
//! - Inference with `no_grad`
//! - Gradient accumulation
//! - Custom differentiable functions
//! - Higher-order gradients
//! - Hardware acceleration
//! - Distributed training
//!
//! Run all examples: `examples::run_all_examples()`
//!
//! # Key Modules
//!
//! ## Core Autograd
//! - [`autograd_traits`]: Core traits for differentiable tensors
//! - [`gradient_storage`]: Thread-safe gradient storage management
//! - [`grad_mode`]: Global gradient computation mode management
//! - [`guards`]: RAII guards for automatic gradient mode restoration
//! - [`variable_env`]: Thread-local variable environment management
//!
//! ## Advanced Features
//! - [`complex_ops`]: Complex number operations with Wirtinger derivatives
//! - [`anomaly_detection`]: Numerical anomaly detection and automatic recovery
//! - [`gradient_clipping`]: Gradient clipping strategies
//! - [`checkpoint_scheduler`]: Memory-efficient gradient checkpointing
//! - [`higher_order_gradients`]: Higher-order derivative computation
//!
//! ## Hardware & Performance
//! - [`hardware_acceleration`]: Multi-platform hardware acceleration
//! - [`profiler`]: Performance profiling and analysis
//! - [`simd_ops`]: SIMD-optimized gradient operations
//! - [`distributed`]: Distributed gradient computation
//!
//! ## Integration & Compatibility
//! - [`pytorch_compat`]: PyTorch compatibility layer
//! - [`jax_transformations`]: JAX-style transformations
//! - [`tensorflow_compat`]: TensorFlow compatibility
//! - [`mlx_compat`]: Apple MLX compatibility
//!
//! # Feature Flags
//!
//! - `default`: Enables autograd, SIMD, and parallel features
//! - `autograd`: SciRS2 autograd integration
//! - `simd`: SIMD optimizations
//! - `parallel`: Parallel gradient computation
//! - `gpu`: GPU acceleration support
//! - `webgpu`: WebGPU for browser deployment
//! - `profiling`: Performance profiling tools
//! - `scirs2-full`: All SciRS2 features

// Core extracted modules for autograd functionality
pub mod anomaly_alerts;
pub mod anomaly_detection;
pub mod autograd_traits;
pub mod complex_ops;
pub mod global_adapter;
pub mod grad_mode;
pub mod gradient_storage;
pub mod guards;
pub mod variable_env;

// Existing specialized modules
pub mod checkpoint_scheduler;
pub mod common_utils;
pub mod communication_efficient;
pub mod compression;
pub mod context;
pub mod differentiable_programming;
pub mod discrete_ops;
pub mod distributed;
pub mod error_diagnostics;
pub mod error_handling;
pub mod external_ad_integration;
pub mod federated_learning;
pub mod flamegraph;
pub mod function;
pub mod function_optimization;
pub mod garbage_collection;
pub mod gradient_checking;
pub mod gradient_filtering;
pub mod gradient_flow_analysis;
pub mod gradient_scaling;
pub mod gradient_scheduler;
pub mod gradient_tracer;
pub mod gradient_validation;
pub mod graph_opt;
pub mod graph_visualization;
pub mod hyperparameter_optimization;
pub mod iterative_solvers;
pub mod jax_transformations;
pub mod matrix_calculus;
pub mod memory;
pub mod meta_gradient;
pub mod metrics_collection;
pub mod mlx_compat;
pub mod onnx_integration;
pub mod operation_introspection;
pub mod operation_replay;
pub mod optimization_diff;
pub mod parameter_server;
pub mod profiler;
pub mod property_testing;
pub mod pytorch_compat;
pub mod scirs2_integration;
pub mod simd_gradient;
pub mod simd_ops;
pub mod staleness_handling;
pub mod stochastic_graphs;
pub mod structured_logging;
pub mod symbolic;
pub mod tensorflow_compat;
pub mod visualization;
pub mod vjp_optimization;

// New modules for enhanced functionality
pub mod auto_tuning;
pub mod automatic_error_recovery;
pub mod blas_integration;
pub mod buffer_optimization;
pub mod cross_framework_verification;
pub mod custom_backends;
pub mod edge_case_handling;
pub mod exception_safety;
pub mod gpu_gradient;
pub mod graceful_degradation;
pub mod gradient_clipping;
pub mod gradient_hooks;
pub mod hardware_acceleration;
pub mod health_diagnostics;
pub mod higher_order_gradients;
pub mod integration_patterns;
pub mod intelligent_chunking;
pub mod interactive_debugger;
pub mod neural_architecture_search;
pub mod neural_ode;
pub mod parallel_gradient;
pub mod performance_regression;
pub mod profiling_debugging_integration;
pub mod progress_reporting;
pub mod quantum_autograd;
pub mod raii_resources;
pub mod regression_testing;
pub mod scirs2_integration_testing;
pub mod specialized_gradient_libs;
pub mod stress_testing;

// Additional framework integration modules
pub mod ad_framework_compatibility;

// API stability and versioning
pub mod stable_api;

// Comprehensive examples and documentation
pub mod examples;

// Production monitoring and observability modules
pub mod audit_logging;
pub mod capacity_planning;
pub mod error_rate_monitoring;
pub mod flamegraph_generation;
pub mod gradient_tracing;
pub mod operation_cost_analysis;
pub mod performance_dashboard;

// Re-exports for convenience
pub use crate::error_handling::{AutogradError, AutogradResult};

pub use crate::autograd_traits::{
    AutogradTensor, AutogradTensorFactory, BackwardTensor, GradientAccumulation,
};

pub use crate::global_adapter::{
    backward_global, create_gradient_tensor, get_global_adapter, get_gradient_global,
};

pub use crate::grad_mode::{
    is_grad_enabled, pop_grad_enabled, push_grad_enabled, set_grad_enabled, with_grad_mode,
};

pub use crate::guards::{enable_grad, no_grad, EnableGradGuard, GradModeGuard, NoGradGuard};

pub use crate::gradient_storage::{
    get_gradient_storage, GlobalGradientStorage, GradientStorage, HashMapGradientStorage,
};

pub use crate::variable_env::{
    clear_variable_env, get_or_create_variable_env, handle_inplace_operation,
    is_variable_env_initialized, validate_inplace_operation, with_variable_env, InplaceConfig,
    InplaceStrategy,
};

pub use crate::complex_ops::backward_complex;
pub use crate::pytorch_compat::backward;

pub use crate::anomaly_detection::{
    detect_complex_anomalies,
    recovery::{
        AnomalyRecoverySystem, RecoveryConfig, RecoveryResult, RecoveryStats, RecoveryStrategy,
    },
};

pub use crate::scirs2_integration::{GradientTensor, SciRS2AutogradAdapter};

pub use crate::auto_tuning::{
    AppliedOptimization, AutoTuningController, OptimizationType, ParameterValue,
    PerformanceSnapshot, TuningConfig, TuningRecommendation, TuningStatistics,
};

pub use crate::error_diagnostics::{
    DiagnosticRecommendation, DiagnosticReport, DiagnosticStatus, DiagnosticsConfig,
    ErrorCorrelation, ErrorDiagnosticsSystem, ErrorPattern, LabeledErrorEvent, MLAnalysisResult,
    MLPatternPrediction, MLPatternRecognitionSystem, MLSystemConfig, PatternLabel, SeverityLevel,
    TemporalContext,
};

// Version information
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
pub const VERSION_MAJOR: u32 = 0;
pub const VERSION_MINOR: u32 = 1;
pub const VERSION_PATCH: u32 = 0;

// Common imports and utilities
use torsh_core::error::{Result, TorshError};

/// Version tracking for tensor operations
///
/// This system tracks tensor versions to detect when in-place operations
/// might invalidate the computation graph.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TensorVersion {
    /// Version number that increments with each modification
    pub version: usize,
    /// Unique tensor identifier
    pub tensor_id: usize,
}

impl TensorVersion {
    /// Create a new tensor version
    pub fn new(tensor_id: usize) -> Self {
        Self {
            version: 0,
            tensor_id,
        }
    }

    /// Increment the version (for in-place operations)
    pub fn increment(&mut self) -> Self {
        self.version += 1;
        *self
    }

    /// Check if this version is compatible with another version
    pub fn is_compatible_with(&self, other: &TensorVersion) -> bool {
        self.tensor_id == other.tensor_id && self.version == other.version
    }
}

/// Global tensor ID counter for unique identification
static TENSOR_ID_COUNTER: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);

/// Generate a unique tensor ID
pub fn new_tensor_id() -> usize {
    TENSOR_ID_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
}

/// In-place operation handling with gradient safety
///
/// This module provides utilities for safely handling in-place tensor operations
/// while preserving gradient computation capabilities.
pub mod inplace_versioning {
    use super::*;

    /// Strategy for handling version conflicts in in-place operations
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum VersionConflictStrategy {
        /// Error on version conflicts
        Error,
        /// Warn about version conflicts but allow operation
        Warn,
        /// Create a copy before in-place operation
        CopyOnWrite,
        /// Silently allow version conflicts
        Allow,
    }

    impl Default for VersionConflictStrategy {
        fn default() -> Self {
            Self::Warn
        }
    }

    /// Check for version conflicts before in-place operations
    pub fn check_version_compatibility(
        current: &TensorVersion,
        expected: &TensorVersion,
        strategy: VersionConflictStrategy,
        operation_name: &str,
    ) -> Result<()> {
        if !current.is_compatible_with(expected) {
            let message = format!(
                "Version conflict in in-place operation '{}': expected version {} but found {}",
                operation_name, expected.version, current.version
            );

            match strategy {
                VersionConflictStrategy::Error => {
                    return Err(TorshError::AutogradError(message));
                }
                VersionConflictStrategy::Warn => {
                    tracing::warn!("{}", message);
                }
                VersionConflictStrategy::CopyOnWrite => {
                    tracing::info!("Triggering copy-on-write for: {}", message);
                }
                VersionConflictStrategy::Allow => {
                    tracing::debug!("Allowing version conflict: {}", message);
                }
            }
        }
        Ok(())
    }

    /// Update tensor version after in-place operation
    pub fn update_version_after_inplace(
        version: &mut TensorVersion,
        operation_name: &str,
    ) -> TensorVersion {
        let old_version = *version;
        let new_version = version.increment();
        tracing::debug!(
            "Updated tensor {} version from {} to {} after in-place operation '{}'",
            version.tensor_id,
            old_version.version,
            new_version.version,
            operation_name
        );
        new_version
    }
}

/// Public prelude for convenient importing
pub mod prelude {
    pub use crate::autograd_traits::{AutogradTensor, BackwardTensor, GradientAccumulation};
    pub use crate::global_adapter::{
        backward_global, create_gradient_tensor, get_global_adapter, get_gradient_global,
    };
    pub use crate::grad_mode::{is_grad_enabled, set_grad_enabled, with_grad_mode};
    pub use crate::gradient_storage::{get_gradient_storage, GradientStorage};
    pub use crate::guards::{enable_grad, no_grad, EnableGradGuard, NoGradGuard};
    pub use crate::variable_env::{with_variable_env, InplaceStrategy};
    pub use crate::{new_tensor_id, TensorVersion};

    // Neural Architecture Search
    pub use crate::neural_architecture_search::{
        ConvOperation, DARTSArchitecture, DARTSPruningResult, IdentityOperation, MixedOperation,
        ProgressiveDARTS, ProgressiveStageInfo, SamplingStrategy, SearchableOperation,
        ZeroOperation, DARTS,
    };

    // Neural ODE
    pub use crate::neural_ode::{
        AdjointMethod, AdjointSolution, IntegrationMethod, NeuralODE, NeuralODELayer, ODESolver,
        ODESolverConfig, ODESystem,
    };

    // Parallel Gradient Computation (SciRS2-Core Integration - Phase 1)
    pub use crate::parallel_gradient::{
        configure_global_parallel, get_global_parallel_computer, get_global_parallel_computer_mut,
        ParallelConfig, ParallelGradientComputer, ParallelStats,
    };

    // GPU Gradient Computation (SciRS2-Core Integration - Phase 2)
    pub use crate::gpu_gradient::{
        get_global_gpu_computer, initialize_global_gpu, is_global_gpu_available, ActivationType,
        GpuBackend, GpuConfig, GpuGradientComputer, GpuStats,
    };

    // SIMD Gradient Computation (SciRS2-Core Integration - Phase 3)
    pub use crate::simd_gradient::{
        configure_global_simd, get_global_simd_computer, get_global_simd_computer_mut,
        SimdCapability, SimdConfig, SimdGradientComputer, SimdStats,
    };

    // Intelligent Chunking (SciRS2-Core Integration - Phase 4)
    pub use crate::intelligent_chunking::{
        configure_global_chunker, get_global_chunker, get_global_chunker_mut, ChunkingConfig,
        ChunkingStats, ChunkingStrategy, IntelligentChunker,
    };

    // Advanced Gradient Clipping
    pub use crate::gradient_clipping::{
        clip_gradients_global, configure_global_clipper, get_global_clipper,
        get_global_clipper_mut, ClippingStats, ClippingStrategy, GradientClipper,
    };

    // Higher-Order Gradients (Hessian, Jacobian)
    pub use crate::higher_order_gradients::{
        configure_global_higher_order, get_global_higher_order, get_global_higher_order_mut,
        ComputationMode, GradientOrder, HigherOrderConfig, HigherOrderGradient, HigherOrderStats,
    };

    // Gradient Hooks System
    pub use crate::gradient_hooks::{
        get_global_hook_manager, get_global_hook_manager_mut, register_global_hook, GradientHook,
        GradientHookManager, HookContext, HookPriority, HookStats, HookType,
    };

    // Automatic Error Recovery
    pub use crate::automatic_error_recovery::{
        get_global_recovery, recover_from_error, AutomaticErrorRecovery, CorrectiveTransform,
        RecoveryAction, RecoveryStrategy, TransientFailureType,
    };
    pub use crate::with_error_recovery;

    // Edge Case Handling
    pub use crate::edge_case_handling::{
        get_global_edge_case_handler, handle_tensor_edge_cases, validate_tensor_shapes,
        EdgeCaseHandler, EdgeCaseStrategy, EdgeCaseTransformation, EdgeCaseType, TensorInfo,
    };

    // Quantum Computing Autograd
    pub use crate::quantum_autograd::{
        Complex, Observable, PauliX, QuantumCircuit, QuantumExpectationValue, QuantumGate,
        QuantumState, QuantumStateGradient, Qubit, RotationY, VQEResult, CNOT, VQE,
    };

    // Cross-Framework Gradient Verification
    pub use crate::cross_framework_verification::{
        get_global_verifier, initialize_verification_frameworks, ComparisonTolerance,
        CrossFrameworkVerifier, FrameworkAdapter, GradientComparisonResult, GradientData,
        MockPyTorchAdapter, SupportedFramework, TorshFrameworkAdapter, VerificationReport,
    };

    // Regression Testing
    pub use crate::regression_testing::{
        get_global_regression_tester, GradientRegressionTester, GradientTestCase,
        RegressionTestResult, RegressionTestStatistics,
    };

    // Exception Safety
    pub use crate::exception_safety::{
        get_global_executor, AutogradTransaction, ComputationGraphGuard, ExceptionSafeExecutor,
        ExceptionSafetyAnalyzer, ExceptionSafetyLevel, GradientStorageGuard, ResourceGuard,
        SafetyViolation, SafetyViolationReport, TransactionOperation,
    };

    // Graceful Degradation
    pub use crate::graceful_degradation::{
        get_global_degradation_manager, DegradationEvent, DegradationStatistics,
        DegradationStrategy, FallbackFunction, GracefulDegradationManager, MatrixMultiplyFallback,
        OperationCategory, UnsupportedOperationInfo, UnsupportedReason,
    };

    // SciRS2 Integration Testing
    pub use crate::scirs2_integration_testing::{
        get_global_integration_tester, run_scirs2_integration_tests, CompatibilitySummary,
        PerformanceSummary, SciRS2IntegrationTestCase, SciRS2IntegrationTestSuite,
        SciRS2IntegrationTester, SciRS2TestResult, SciRS2Version, TestCategory,
        TestCategoryResults,
    };

    // Integration Patterns and Documentation
    pub use crate::integration_patterns::{
        IntegrationDocumentation, IntegrationPatterns, MigrationGuide, MigrationScenario, Pattern,
        PatternCategory, PatternDocumentation, TroubleshootingGuide, TroubleshootingIssue,
    };

    // BLAS Integration
    pub use crate::blas_integration::{
        blas_dot, blas_gemm, blas_gemv, get_global_blas_manager, BlasConfig, BlasImplementation,
        BlasManager, BlasOperation, BlasPerformanceReport, BlasProvider, PureRustBlasProvider,
    };

    // Specialized Gradient Libraries
    pub use crate::specialized_gradient_libs::{
        get_global_specialized_manager, BenchmarkReport, CasADiLibrary, ComputationResult,
        Function, GradientComputationType, LibraryUsageReport, QuadraticFunction, SparseGradient,
        SpecializedGradientLibrary, SpecializedLibConfig, SpecializedLibrary,
        SpecializedLibraryManager,
    };

    // Custom Autograd Backends
    pub use crate::custom_backends::{
        get_active_backend, get_global_backend_registry, AutogradBackend, BackendCapability,
        BackendConfig, BackendInfo, BackendRegistry, BackendTensor, CustomOperation, DataType,
        DeviceConfig, DeviceType, GradFunction, OperationContext, OptimizationLevel,
        PerformanceStats, ReferenceBackend,
    };

    // Hardware Acceleration
    pub use crate::hardware_acceleration::{
        get_global_acceleration_manager, AccelerationConfig, AcceleratorBenchmarkReport,
        AcceleratorType, AcceleratorUsageReport, Conv2DParams, CudaAccelerator, DeviceStats,
        HardwareAccelerationManager, HardwareAccelerator, HardwareCapability, HardwareDevice,
        HardwareMemoryHandle, MetalAccelerator, OptimizationLevel as HardwareOptimizationLevel,
        PrecisionPreference,
    };

    // Profiling and Debugging Integration
    pub use crate::profiling_debugging_integration::{
        get_global_profiling_debugging_manager, AnalysisCapability, CPUProfile, DebuggingConfig,
        DebuggingReport, DebuggingSession, DebuggingTool, ExternalDebugger, ExternalProfiler,
        GPUProfile, GdbDebugger, Hotspot, IntegrationConfig, IntegrationReport, MemoryError,
        MemoryProfile, PerfProfiler, ProfilingConfig, ProfilingDebuggingManager, ProfilingReport,
        ProfilingSession, ProfilingTool, StackTrace, ThreadError,
    };

    // AD Framework Compatibility
    pub use crate::ad_framework_compatibility::{
        check_framework_compatibility, convert_tensor, get_global_compatibility_manager,
        migrate_model, ADFramework, ADFrameworkCompatibilityManager, AutomationLevel,
        CompatibilityLevel, CompatibilityReport, CustomOperationDefinition, EffortLevel,
        FrameworkAdapter as ADFrameworkAdapter, FrameworkCapabilities, FrameworkTensor,
        MigrationCapability, MigrationData, MigrationOperation, MigrationPlan, MigrationResult,
        MigrationStep, PerformanceComparison, PerformanceMetrics, PyTorchAdapter, PyTorchTensor,
        RequiredTransformation, UniversalDataType, UniversalOperation, UniversalTensor,
        ValidationResult,
    };
}

/// Accumulate gradients with overflow protection
pub mod accumulate {
    use super::*;
    use crate::autograd_traits::AutogradTensor;
    use num_traits::Float;

    /// Accumulate gradients safely with overflow detection
    pub fn accumulate_gradient_safe<T>(
        existing: &mut dyn AutogradTensor<T>,
        new_grad: &dyn AutogradTensor<T>,
        overflow_threshold: Option<T>,
    ) -> Result<()>
    where
        T: torsh_core::dtype::TensorElement
            + Float
            + Clone
            + std::fmt::Debug
            + std::fmt::Display
            + Send
            + Sync,
    {
        // Check shapes match
        if existing.shape() != new_grad.shape() {
            return Err(TorshError::AutogradError(format!(
                "Shape mismatch in gradient accumulation: {:?} vs {:?}",
                existing.shape(),
                new_grad.shape()
            )));
        }

        // Get data for accumulation
        let existing_data = existing.data();
        let new_data = new_grad.data();

        // Check for overflow if threshold provided
        if let Some(threshold) = overflow_threshold {
            for (existing_val, new_val) in existing_data.iter().zip(new_data.iter()) {
                let sum = *existing_val + *new_val;
                if sum.abs() > threshold {
                    return Err(TorshError::AutogradError(format!(
                        "Gradient accumulation overflow detected: {} + {} = {} > {}",
                        existing_val, new_val, sum, threshold
                    )));
                }
            }
        }

        // Since we can't modify trait objects directly, this is a limitation
        // Real implementation would be in concrete tensor types
        tracing::debug!(
            "Gradient accumulation requested for tensor with shape {:?}",
            existing.shape()
        );

        Ok(())
    }

    /// Check if gradient accumulation would cause overflow
    pub fn check_accumulation_overflow<T>(val1: T, val2: T, threshold: T) -> bool
    where
        T: Float + PartialOrd,
    {
        let sum = val1 + val2;
        sum.abs() > threshold
    }
}

/// Gradient clipping utilities
pub mod clip {
    use super::*;
    use crate::autograd_traits::AutogradTensor;
    use num_traits::Float;

    /// Clip gradients by global norm
    pub fn clip_grad_norm<T>(
        gradients: &[&dyn AutogradTensor<T>],
        max_norm: T,
        norm_type: f32,
    ) -> Result<T>
    where
        T: torsh_core::dtype::TensorElement
            + Float
            + Clone
            + std::fmt::Debug
            + Send
            + Sync
            + std::fmt::Display,
        f32: From<T>,
    {
        if gradients.is_empty() {
            return Ok(<T as num_traits::Zero>::zero());
        }

        // Calculate total norm
        let mut total_norm = <T as num_traits::Zero>::zero();

        for grad in gradients {
            let data = grad.data();
            for &val in data.iter() {
                if norm_type == 2.0 {
                    total_norm = total_norm + val * val;
                } else if norm_type == 1.0 {
                    total_norm = total_norm + val.abs();
                } else {
                    let abs_val = val.abs();
                    total_norm = total_norm
                        + abs_val.powf(T::from(norm_type).unwrap_or(<T as num_traits::One>::one()));
                }
            }
        }

        if norm_type == 2.0 {
            total_norm = total_norm.sqrt();
        } else if norm_type != 1.0 {
            total_norm =
                total_norm.powf(T::from(1.0 / norm_type).unwrap_or(<T as num_traits::One>::one()));
        }

        tracing::debug!("Calculated gradient norm: {:?}", total_norm);

        // Calculate clipping ratio
        let clip_coef =
            max_norm / (total_norm + T::from(1e-6).unwrap_or(<T as num_traits::One>::one()));
        let clip_coef = clip_coef.min(<T as num_traits::One>::one());

        tracing::debug!("Gradient clipping coefficient: {:?}", clip_coef);

        Ok(total_norm)
    }

    /// Clip gradients by value
    pub fn clip_grad_value<T>(
        gradient: &dyn AutogradTensor<T>,
        min_value: T,
        max_value: T,
    ) -> Result<Vec<T>>
    where
        T: torsh_core::dtype::TensorElement + Float + Clone + std::fmt::Debug + Send + Sync,
    {
        let data = gradient.data();
        let clipped: Vec<T> = data
            .iter()
            .map(|&val| val.max(min_value).min(max_value))
            .collect();

        Ok(clipped)
    }
}

/// Forward-mode automatic differentiation
pub mod forward_mode {

    use num_traits::Float;

    /// Dual number for forward-mode AD
    #[derive(Debug, Clone, Copy, PartialEq)]
    pub struct Dual<T> {
        /// Value
        pub value: T,
        /// Derivative
        pub derivative: T,
    }

    impl<T: Float> Dual<T> {
        /// Create a new dual number
        pub fn new(value: T, derivative: T) -> Self {
            Self { value, derivative }
        }

        /// Create a variable (derivative = 1)
        pub fn variable(value: T) -> Self {
            Self::new(value, <T as num_traits::One>::one())
        }

        /// Create a constant (derivative = 0)
        pub fn constant(value: T) -> Self {
            Self::new(value, T::zero())
        }
    }

    impl<T: Float> std::ops::Add for Dual<T> {
        type Output = Self;

        fn add(self, rhs: Self) -> Self::Output {
            Self::new(self.value + rhs.value, self.derivative + rhs.derivative)
        }
    }

    impl<T: Float> std::ops::Mul for Dual<T> {
        type Output = Self;

        fn mul(self, rhs: Self) -> Self::Output {
            Self::new(
                self.value * rhs.value,
                self.derivative * rhs.value + self.value * rhs.derivative,
            )
        }
    }

    /// Compute forward-mode derivative
    pub fn forward_derivative<T, F>(input: T, f: F) -> (T, T)
    where
        T: Float + Clone,
        F: Fn(Dual<T>) -> Dual<T>,
    {
        let dual_input = Dual::variable(input);
        let dual_output = f(dual_input);
        (dual_output.value, dual_output.derivative)
    }
}

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

    #[test]
    fn test_tensor_version() {
        let mut version = TensorVersion::new(42);
        assert_eq!(version.tensor_id, 42);
        assert_eq!(version.version, 0);

        let new_version = version.increment();
        assert_eq!(new_version.version, 1);
        assert_eq!(version.version, 1);
    }

    #[test]
    fn test_tensor_version_compatibility() {
        let version1 = TensorVersion::new(1);
        let version2 = TensorVersion::new(1);
        let version3 = TensorVersion::new(2);

        assert!(version1.is_compatible_with(&version2));
        assert!(!version1.is_compatible_with(&version3));
    }

    #[test]
    fn test_unique_tensor_ids() {
        let id1 = new_tensor_id();
        let id2 = new_tensor_id();
        let id3 = new_tensor_id();

        assert_ne!(id1, id2);
        assert_ne!(id2, id3);
        assert_ne!(id1, id3);
    }

    #[test]
    fn test_dual_number_arithmetic() {
        use forward_mode::Dual;

        let x = Dual::new(2.0, 1.0);
        let y = Dual::new(3.0, 0.0);

        let sum = x + y;
        assert_eq!(sum.value, 5.0);
        assert_eq!(sum.derivative, 1.0);

        let product = x * y;
        assert_eq!(product.value, 6.0);
        assert_eq!(product.derivative, 3.0);
    }

    #[test]
    fn test_forward_derivative() {
        use forward_mode::forward_derivative;

        // f(x) = x^2, f'(x) = 2x
        let (value, derivative) = forward_derivative(3.0, |x| x * x);
        assert_eq!(value, 9.0);
        assert_eq!(derivative, 6.0);
    }
}