sklears_neural/

gradient_checking.rs

use scirs2_core::ndarray::{Array1, Array2, Axis, ScalarOperand};
use scirs2_core::numeric::{Float, One, ToPrimitive};
use std::fmt::Debug;

use crate::activation::Activation;
use crate::self_supervised::{DenseLayer, SimpleMLP};
use sklears_core::error::SklearsError;
use sklears_core::types::FloatBounds;

/// Numerical gradient checking utilities for neural networks
///
/// This module provides tools for validating analytical gradients by comparing them
/// with numerically computed gradients using finite differences.

/// Gradient checking configuration
#[derive(Debug, Clone)]
pub struct GradientCheckConfig<T: Float> {
    /// Finite difference step size (epsilon)
    pub epsilon: T,
    /// Relative tolerance for gradient comparison
    pub relative_tolerance: T,
    /// Absolute tolerance for gradient comparison
    pub absolute_tolerance: T,
    /// Whether to use centered differences (more accurate but 2x slower)
    pub use_centered_differences: bool,
    /// Maximum number of parameters to check (for efficiency)
    pub max_params_to_check: Option<usize>,
    /// Random seed for parameter sampling
    pub random_seed: Option<u64>,
}

impl<T: Float> Default for GradientCheckConfig<T> {
    fn default() -> Self {
        Self {
            epsilon: T::from(1e-7).unwrap(),
            relative_tolerance: T::from(1e-5).unwrap(),
            absolute_tolerance: T::from(1e-8).unwrap(),
            use_centered_differences: true,
            max_params_to_check: Some(100),
            random_seed: Some(42),
        }
    }
}

/// Results of gradient checking
#[derive(Debug, Clone)]
pub struct GradientCheckResults<T: Float> {
    /// Whether all gradients passed the check
    pub all_passed: bool,
    /// Number of parameters checked
    pub num_checked: usize,
    /// Number of parameters that passed
    pub num_passed: usize,
    /// Maximum relative error found
    pub max_relative_error: T,
    /// Maximum absolute error found
    pub max_absolute_error: T,
    /// Average relative error
    pub avg_relative_error: T,
    /// Average absolute error
    pub avg_absolute_error: T,
    /// Detailed results for each parameter
    pub parameter_results: Vec<ParameterGradientResult<T>>,
}

/// Results for a single parameter gradient check
#[derive(Debug, Clone)]
pub struct ParameterGradientResult<T: Float> {
    /// Parameter index/identifier
    pub param_index: usize,
    /// Analytical gradient value
    pub analytical_gradient: T,
    /// Numerical gradient value
    pub numerical_gradient: T,
    /// Relative error
    pub relative_error: T,
    /// Absolute error
    pub absolute_error: T,
    /// Whether this parameter passed the check
    pub passed: bool,
}

/// Loss function trait for gradient checking
pub trait LossFunction<T: FloatBounds + ScalarOperand> {
    /// Compute loss given predictions and targets
    fn compute_loss(&self, predictions: &Array2<T>, targets: &Array2<T>)
        -> Result<T, SklearsError>;

    /// Compute loss gradient with respect to predictions
    fn compute_gradient(
        &self,
        predictions: &Array2<T>,
        targets: &Array2<T>,
    ) -> Result<Array2<T>, SklearsError>;
}

/// Mean Squared Error loss function
#[derive(Debug, Clone)]
pub struct MeanSquaredError<T: FloatBounds + ScalarOperand> {
    _phantom: std::marker::PhantomData<T>,
}

impl<T: FloatBounds + ScalarOperand> MeanSquaredError<T> {
    pub fn new() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<T: FloatBounds + ScalarOperand> LossFunction<T> for MeanSquaredError<T> {
    fn compute_loss(
        &self,
        predictions: &Array2<T>,
        targets: &Array2<T>,
    ) -> Result<T, SklearsError> {
        let diff = predictions - targets;
        let squared_diff = diff.mapv(|x| x * x);
        let mse = squared_diff.sum() / T::from(predictions.len()).unwrap();
        Ok(mse)
    }

    fn compute_gradient(
        &self,
        predictions: &Array2<T>,
        targets: &Array2<T>,
    ) -> Result<Array2<T>, SklearsError> {
        let diff = predictions - targets;
        let factor = T::from(2.0).unwrap() / T::from(predictions.len()).unwrap();
        Ok(diff * factor)
    }
}

/// Cross-entropy loss function
#[derive(Debug, Clone)]
pub struct CrossEntropyLoss<T: FloatBounds + ScalarOperand> {
    _phantom: std::marker::PhantomData<T>,
}

impl<T: FloatBounds + ScalarOperand> CrossEntropyLoss<T> {
    pub fn new() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<T: FloatBounds + ScalarOperand> LossFunction<T> for CrossEntropyLoss<T> {
    fn compute_loss(
        &self,
        predictions: &Array2<T>,
        targets: &Array2<T>,
    ) -> Result<T, SklearsError> {
        let epsilon = T::from(1e-15).unwrap();
        let clipped_preds = predictions.mapv(|x| x.max(epsilon).min(T::one() - epsilon));

        let log_preds = clipped_preds.mapv(|x| x.ln());
        let loss = -(targets * log_preds).sum() / T::from(predictions.nrows()).unwrap();
        Ok(loss)
    }

    fn compute_gradient(
        &self,
        predictions: &Array2<T>,
        targets: &Array2<T>,
    ) -> Result<Array2<T>, SklearsError> {
        let epsilon = T::from(1e-15).unwrap();
        let clipped_preds = predictions.mapv(|x| x.max(epsilon).min(T::one() - epsilon));

        let grad = -(targets / clipped_preds) / T::from(predictions.nrows()).unwrap();
        Ok(grad)
    }
}

/// Gradient checker for neural networks
#[derive(Debug)]
pub struct GradientChecker<T: FloatBounds + ScalarOperand + ToPrimitive> {
    config: GradientCheckConfig<T>,
}

impl<T: FloatBounds + ScalarOperand + ToPrimitive> GradientChecker<T> {
    /// Create a new gradient checker
    pub fn new(config: GradientCheckConfig<T>) -> Self {
        Self { config }
    }

    /// Check gradients for a neural network
    pub fn check_network_gradients(
        &self,
        network: &mut SimpleMLP<T>,
        inputs: &Array2<T>,
        targets: &Array2<T>,
        loss_fn: &dyn LossFunction<T>,
    ) -> Result<GradientCheckResults<T>, SklearsError> {
        // Forward pass to get predictions
        let predictions = network.forward(inputs)?;

        // Compute analytical gradients using backpropagation
        let analytical_grads =
            self.compute_analytical_gradients(network, inputs, targets, loss_fn)?;

        // Compute numerical gradients using finite differences
        let numerical_grads =
            self.compute_numerical_gradients(network, inputs, targets, loss_fn)?;

        // Compare gradients and generate results
        self.compare_gradients(&analytical_grads, &numerical_grads)
    }

    /// Check gradients for a single layer
    pub fn check_layer_gradients(
        &self,
        layer: &mut DenseLayer<T>,
        inputs: &Array2<T>,
        output_gradients: &Array2<T>,
    ) -> Result<GradientCheckResults<T>, SklearsError> {
        // This is a simplified version for demonstration
        // In practice, you'd implement full gradient checking for each layer type

        let mut parameter_results = Vec::new();
        let mut num_passed = 0;
        let mut max_rel_error = T::zero();
        let mut max_abs_error = T::zero();
        let mut sum_rel_error = T::zero();
        let mut sum_abs_error = T::zero();

        // For demonstration, we'll just check a few parameters
        // In practice, you'd check all weights and biases
        let num_to_check = std::cmp::min(10, 100); // Simplified

        for i in 0..num_to_check {
            let analytical_grad = T::from(0.1).unwrap(); // Placeholder
            let numerical_grad = T::from(0.101).unwrap(); // Placeholder

            let abs_error = (analytical_grad - numerical_grad).abs();
            let rel_error = if numerical_grad.abs() > T::zero() {
                abs_error / numerical_grad.abs()
            } else {
                abs_error
            };

            let passed = rel_error < self.config.relative_tolerance
                && abs_error < self.config.absolute_tolerance;

            if passed {
                num_passed += 1;
            }

            max_rel_error = max_rel_error.max(rel_error);
            max_abs_error = max_abs_error.max(abs_error);
            sum_rel_error = sum_rel_error + rel_error;
            sum_abs_error = sum_abs_error + abs_error;

            parameter_results.push(ParameterGradientResult {
                param_index: i,
                analytical_gradient: analytical_grad,
                numerical_gradient: numerical_grad,
                relative_error: rel_error,
                absolute_error: abs_error,
                passed,
            });
        }

        let avg_rel_error = sum_rel_error / T::from(num_to_check).unwrap();
        let avg_abs_error = sum_abs_error / T::from(num_to_check).unwrap();

        Ok(GradientCheckResults {
            all_passed: num_passed == num_to_check,
            num_checked: num_to_check,
            num_passed,
            max_relative_error: max_rel_error,
            max_absolute_error: max_abs_error,
            avg_relative_error: avg_rel_error,
            avg_absolute_error: avg_abs_error,
            parameter_results,
        })
    }

    /// Compute analytical gradients using backpropagation
    fn compute_analytical_gradients(
        &self,
        network: &mut SimpleMLP<T>,
        inputs: &Array2<T>,
        targets: &Array2<T>,
        loss_fn: &dyn LossFunction<T>,
    ) -> Result<Vec<Array1<T>>, SklearsError> {
        // Forward pass
        let predictions = network.forward(inputs)?;

        // Compute loss gradient
        let loss_grad = loss_fn.compute_gradient(&predictions, targets)?;

        // Backward pass through network
        // This is simplified - in practice you'd implement full backpropagation
        let mut gradients = Vec::new();

        // For demonstration, we'll create some dummy gradients
        // In practice, this would be the actual backpropagation implementation
        for i in 0..10 {
            let grad = Array1::from_vec(vec![T::from(i as f64 * 0.1).unwrap(); 10]);
            gradients.push(grad);
        }

        Ok(gradients)
    }

    /// Compute numerical gradients using finite differences
    fn compute_numerical_gradients(
        &self,
        network: &mut SimpleMLP<T>,
        inputs: &Array2<T>,
        targets: &Array2<T>,
        loss_fn: &dyn LossFunction<T>,
    ) -> Result<Vec<Array1<T>>, SklearsError> {
        let mut numerical_grads = Vec::new();

        // For each parameter in the network
        for param_group in 0..10 {
            // Simplified
            let mut param_grads = Vec::new();

            for param_idx in 0..10 {
                // Simplified
                let grad = if self.config.use_centered_differences {
                    self.compute_centered_difference(
                        network,
                        inputs,
                        targets,
                        loss_fn,
                        param_group,
                        param_idx,
                    )?
                } else {
                    self.compute_forward_difference(
                        network,
                        inputs,
                        targets,
                        loss_fn,
                        param_group,
                        param_idx,
                    )?
                };
                param_grads.push(grad);
            }

            numerical_grads.push(Array1::from_vec(param_grads));
        }

        Ok(numerical_grads)
    }

    /// Compute centered finite difference for a parameter
    fn compute_centered_difference(
        &self,
        network: &mut SimpleMLP<T>,
        inputs: &Array2<T>,
        targets: &Array2<T>,
        loss_fn: &dyn LossFunction<T>,
        param_group: usize,
        param_idx: usize,
    ) -> Result<T, SklearsError> {
        // Get current parameter value
        let original_param = T::from(0.5).unwrap(); // Placeholder

        // Compute loss with parameter + epsilon
        // (This is simplified - in practice you'd modify the actual network parameters)
        let loss_plus = self.compute_loss_with_perturbed_param(
            network,
            inputs,
            targets,
            loss_fn,
            param_group,
            param_idx,
            original_param + self.config.epsilon,
        )?;

        // Compute loss with parameter - epsilon
        let loss_minus = self.compute_loss_with_perturbed_param(
            network,
            inputs,
            targets,
            loss_fn,
            param_group,
            param_idx,
            original_param - self.config.epsilon,
        )?;

        // Centered difference
        let grad = (loss_plus - loss_minus) / (T::from(2.0).unwrap() * self.config.epsilon);
        Ok(grad)
    }

    /// Compute forward finite difference for a parameter
    fn compute_forward_difference(
        &self,
        network: &mut SimpleMLP<T>,
        inputs: &Array2<T>,
        targets: &Array2<T>,
        loss_fn: &dyn LossFunction<T>,
        param_group: usize,
        param_idx: usize,
    ) -> Result<T, SklearsError> {
        // Get current parameter value
        let original_param = T::from(0.5).unwrap(); // Placeholder

        // Compute original loss
        let original_loss = self.compute_loss_with_perturbed_param(
            network,
            inputs,
            targets,
            loss_fn,
            param_group,
            param_idx,
            original_param,
        )?;

        // Compute loss with parameter + epsilon
        let perturbed_loss = self.compute_loss_with_perturbed_param(
            network,
            inputs,
            targets,
            loss_fn,
            param_group,
            param_idx,
            original_param + self.config.epsilon,
        )?;

        // Forward difference
        let grad = (perturbed_loss - original_loss) / self.config.epsilon;
        Ok(grad)
    }

    /// Compute loss with a perturbed parameter (simplified)
    fn compute_loss_with_perturbed_param(
        &self,
        network: &mut SimpleMLP<T>,
        inputs: &Array2<T>,
        targets: &Array2<T>,
        loss_fn: &dyn LossFunction<T>,
        _param_group: usize,
        _param_idx: usize,
        _param_value: T,
    ) -> Result<T, SklearsError> {
        // This is simplified - in practice you'd:
        // 1. Save the original parameter value
        // 2. Set the parameter to the new value
        // 3. Run forward pass
        // 4. Compute loss
        // 5. Restore original parameter value

        let predictions = network.forward(inputs)?;
        loss_fn.compute_loss(&predictions, targets)
    }

    /// Compare analytical and numerical gradients
    fn compare_gradients(
        &self,
        analytical: &[Array1<T>],
        numerical: &[Array1<T>],
    ) -> Result<GradientCheckResults<T>, SklearsError> {
        let mut parameter_results = Vec::new();
        let mut num_passed = 0;
        let mut max_rel_error = T::zero();
        let mut max_abs_error = T::zero();
        let mut sum_rel_error = T::zero();
        let mut sum_abs_error = T::zero();
        let mut total_checked = 0;

        for (group_idx, (anal_group, num_group)) in
            analytical.iter().zip(numerical.iter()).enumerate()
        {
            for (param_idx, (&anal_grad, &num_grad)) in
                anal_group.iter().zip(num_group.iter()).enumerate()
            {
                let abs_error = (anal_grad - num_grad).abs();
                let rel_error = if num_grad.abs() > T::zero() {
                    abs_error / num_grad.abs()
                } else {
                    abs_error
                };

                let passed = rel_error < self.config.relative_tolerance
                    && abs_error < self.config.absolute_tolerance;

                if passed {
                    num_passed += 1;
                }

                max_rel_error = max_rel_error.max(rel_error);
                max_abs_error = max_abs_error.max(abs_error);
                sum_rel_error = sum_rel_error + rel_error;
                sum_abs_error = sum_abs_error + abs_error;
                total_checked += 1;

                parameter_results.push(ParameterGradientResult {
                    param_index: group_idx * 1000 + param_idx, // Simple encoding
                    analytical_gradient: anal_grad,
                    numerical_gradient: num_grad,
                    relative_error: rel_error,
                    absolute_error: abs_error,
                    passed,
                });

                // Limit number of parameters checked for efficiency
                if let Some(max_params) = self.config.max_params_to_check {
                    if total_checked >= max_params {
                        break;
                    }
                }
            }

            if let Some(max_params) = self.config.max_params_to_check {
                if total_checked >= max_params {
                    break;
                }
            }
        }

        let avg_rel_error = if total_checked > 0 {
            sum_rel_error / T::from(total_checked).unwrap()
        } else {
            T::zero()
        };

        let avg_abs_error = if total_checked > 0 {
            sum_abs_error / T::from(total_checked).unwrap()
        } else {
            T::zero()
        };

        Ok(GradientCheckResults {
            all_passed: num_passed == total_checked,
            num_checked: total_checked,
            num_passed,
            max_relative_error: max_rel_error,
            max_absolute_error: max_abs_error,
            avg_relative_error: avg_rel_error,
            avg_absolute_error: avg_abs_error,
            parameter_results,
        })
    }
}

/// Utility functions for gradient checking
impl<T: FloatBounds + ScalarOperand + ToPrimitive> GradientChecker<T> {
    /// Check if gradients are approximately equal
    pub fn gradients_are_equal(&self, analytical: T, numerical: T) -> bool {
        let abs_error = (analytical - numerical).abs();
        let rel_error = if numerical.abs() > T::zero() {
            abs_error / numerical.abs()
        } else {
            abs_error
        };

        rel_error < self.config.relative_tolerance && abs_error < self.config.absolute_tolerance
    }

    /// Compute relative error between two gradients
    pub fn compute_relative_error(&self, analytical: T, numerical: T) -> T {
        let abs_error = (analytical - numerical).abs();
        if numerical.abs() > T::zero() {
            abs_error / numerical.abs()
        } else {
            abs_error
        }
    }

    /// Generate a summary report of gradient checking results
    pub fn generate_report(&self, results: &GradientCheckResults<T>) -> String {
        let mut report = String::new();

        report.push_str("=== Gradient Checking Report ===\n");
        report.push_str(&format!(
            "Overall Status: {}\n",
            if results.all_passed {
                "PASSED"
            } else {
                "FAILED"
            }
        ));
        report.push_str(&format!("Parameters Checked: {}\n", results.num_checked));
        report.push_str(&format!("Parameters Passed: {}\n", results.num_passed));
        report.push_str(&format!(
            "Pass Rate: {:.2}%\n",
            (results.num_passed as f64 / results.num_checked as f64) * 100.0
        ));
        report.push_str(&format!(
            "Max Relative Error: {:.2e}\n",
            results.max_relative_error.to_f64().unwrap_or(0.0)
        ));
        report.push_str(&format!(
            "Max Absolute Error: {:.2e}\n",
            results.max_absolute_error.to_f64().unwrap_or(0.0)
        ));
        report.push_str(&format!(
            "Avg Relative Error: {:.2e}\n",
            results.avg_relative_error.to_f64().unwrap_or(0.0)
        ));
        report.push_str(&format!(
            "Avg Absolute Error: {:.2e}\n",
            results.avg_absolute_error.to_f64().unwrap_or(0.0)
        ));

        // Add details for failed parameters
        let failed_params: Vec<_> = results
            .parameter_results
            .iter()
            .filter(|r| !r.passed)
            .collect();

        if !failed_params.is_empty() {
            report.push_str("\nFailed Parameters:\n");
            for param in failed_params.iter().take(10) {
                // Show first 10 failures
                report.push_str(&format!(
                    "  Param {}: analytical={:.6e}, numerical={:.6e}, rel_err={:.2e}, abs_err={:.2e}\n",
                    param.param_index,
                    param.analytical_gradient.to_f64().unwrap_or(0.0),
                    param.numerical_gradient.to_f64().unwrap_or(0.0),
                    param.relative_error.to_f64().unwrap_or(0.0),
                    param.absolute_error.to_f64().unwrap_or(0.0)
                ));
            }

            if failed_params.len() > 10 {
                report.push_str(&format!(
                    "  ... and {} more failures\n",
                    failed_params.len() - 10
                ));
            }
        }

        report
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;

    #[test]
    fn test_gradient_check_config_default() {
        let config = GradientCheckConfig::<f32>::default();
        assert!(config.epsilon > 0.0);
        assert!(config.use_centered_differences);
        assert_eq!(config.max_params_to_check, Some(100));
    }

    #[test]
    fn test_mse_loss_function() {
        let mse = MeanSquaredError::<f32>::new();

        let predictions = Array2::from_shape_vec((2, 2), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let targets = Array2::from_shape_vec((2, 2), vec![1.1, 1.9, 3.1, 3.9]).unwrap();

        let loss = mse.compute_loss(&predictions, &targets).unwrap();
        assert!(loss > 0.0);

        let gradient = mse.compute_gradient(&predictions, &targets).unwrap();
        assert_eq!(gradient.dim(), predictions.dim());
    }

    #[test]
    fn test_cross_entropy_loss_function() {
        let ce = CrossEntropyLoss::<f32>::new();

        let predictions = Array2::from_shape_vec((2, 2), vec![0.8, 0.2, 0.3, 0.7]).unwrap();
        let targets = Array2::from_shape_vec((2, 2), vec![1.0, 0.0, 0.0, 1.0]).unwrap();

        let loss = ce.compute_loss(&predictions, &targets).unwrap();
        assert!(loss > 0.0);

        let gradient = ce.compute_gradient(&predictions, &targets).unwrap();
        assert_eq!(gradient.dim(), predictions.dim());
    }

    #[test]
    fn test_gradient_checker_creation() {
        let config = GradientCheckConfig::<f32>::default();
        let checker = GradientChecker::new(config);
        assert!(checker.config.epsilon > 0.0);
    }

    #[test]
    fn test_gradients_are_equal() {
        let config = GradientCheckConfig {
            epsilon: 1e-7,
            relative_tolerance: 1e-5,
            absolute_tolerance: 1e-6, // Increased tolerance for the test
            use_centered_differences: true,
            max_params_to_check: Some(100),
            random_seed: Some(42),
        };
        let checker = GradientChecker::new(config);

        // Test equal gradients
        assert!(checker.gradients_are_equal(1.0, 1.0));

        // Test nearly equal gradients (within tolerance)
        assert!(checker.gradients_are_equal(1.0, 1.000001));

        // Test different gradients
        assert!(!checker.gradients_are_equal(1.0, 1.1));
    }

    #[test]
    fn test_compute_relative_error() {
        let config = GradientCheckConfig::<f32>::default();
        let checker = GradientChecker::new(config);

        let rel_error = checker.compute_relative_error(1.0, 1.1);
        assert_abs_diff_eq!(rel_error, 0.090909, epsilon = 1e-5);

        // Test with zero numerical gradient
        let rel_error_zero = checker.compute_relative_error(0.1, 0.0);
        assert_abs_diff_eq!(rel_error_zero, 0.1, epsilon = 1e-6);
    }

    #[test]
    fn test_parameter_gradient_result() {
        let result = ParameterGradientResult {
            param_index: 0,
            analytical_gradient: 1.0,
            numerical_gradient: 1.01,
            relative_error: 0.0099,
            absolute_error: 0.01,
            passed: true,
        };

        assert_eq!(result.param_index, 0);
        assert!(result.passed);
        assert_eq!(result.analytical_gradient, 1.0);
    }

    #[test]
    fn test_gradient_check_results() {
        let param_results = vec![
            ParameterGradientResult {
                param_index: 0,
                analytical_gradient: 1.0,
                numerical_gradient: 1.01,
                relative_error: 0.0099,
                absolute_error: 0.01,
                passed: true,
            },
            ParameterGradientResult {
                param_index: 1,
                analytical_gradient: 2.0,
                numerical_gradient: 2.2,
                relative_error: 0.091,
                absolute_error: 0.2,
                passed: false,
            },
        ];

        let results = GradientCheckResults {
            all_passed: false,
            num_checked: 2,
            num_passed: 1,
            max_relative_error: 0.091,
            max_absolute_error: 0.2,
            avg_relative_error: 0.05045,
            avg_absolute_error: 0.105,
            parameter_results: param_results,
        };

        assert!(!results.all_passed);
        assert_eq!(results.num_checked, 2);
        assert_eq!(results.num_passed, 1);
    }

    #[test]
    fn test_generate_report() {
        let config = GradientCheckConfig::<f32>::default();
        let checker = GradientChecker::new(config);

        let results = GradientCheckResults {
            all_passed: true,
            num_checked: 10,
            num_passed: 10,
            max_relative_error: 1e-6,
            max_absolute_error: 1e-8,
            avg_relative_error: 1e-7,
            avg_absolute_error: 1e-9,
            parameter_results: Vec::new(),
        };

        let report = checker.generate_report(&results);
        assert!(report.contains("PASSED"));
        assert!(report.contains("Parameters Checked: 10"));
        assert!(report.contains("Pass Rate: 100.00%"));
    }

    #[test]
    fn test_layer_gradient_checking() {
        let config = GradientCheckConfig::<f32>::default();
        let checker = GradientChecker::new(config);

        let mut layer = DenseLayer::<f32>::new(5, 3, Some(Activation::Relu));
        let inputs = Array2::from_shape_vec((2, 5), vec![1.0; 10]).unwrap();
        let output_grads = Array2::from_shape_vec((2, 3), vec![0.1; 6]).unwrap();

        let results = checker
            .check_layer_gradients(&mut layer, &inputs, &output_grads)
            .unwrap();
        assert!(results.num_checked > 0);
        // Note: This is a simplified test since the actual gradient checking is not fully implemented
    }
}