torsh_nn/core/

module_ext.rs

//! Module trait ergonomic extensions
//!
//! This module provides additional ergonomic helpers and utilities for the Module trait,
//! following Rust best practices for trait extension patterns.

use crate::Module;
use torsh_core::device::DeviceType;
use torsh_core::error::Result;
use torsh_tensor::Tensor;

#[cfg(feature = "std")]
use std::collections::HashMap;

#[cfg(not(feature = "std"))]
use hashbrown::HashMap;

/// Extension trait providing additional ergonomic methods for Module
///
/// This trait is automatically implemented for all types that implement Module,
/// providing additional convenience methods without requiring changes to existing code.
///
/// # Design Philosophy
///
/// This extension follows Rust's extension trait pattern to:
/// - Add functionality without breaking backward compatibility
/// - Keep the core Module trait focused on essential methods
/// - Provide advanced features for users who need them
/// - Enable fluent/builder-style APIs
pub trait ModuleExt: Module {
    // === Fluent API / Builder Pattern Methods ===

    /// Chain forward pass with a transformation function
    ///
    /// This enables functional-style chaining of operations.
    ///
    /// # Arguments
    /// * `input` - Input tensor
    /// * `f` - Transformation function to apply to output
    ///
    /// # Returns
    /// * `Result<Tensor>` - Transformed output
    ///
    /// # Example
    /// ```ignore
    /// let output = layer.and_then(&input, |x| x.relu())?;
    /// ```
    fn and_then<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
    where
        F: FnOnce(Tensor) -> Result<Tensor>,
    {
        let output = self.forward(input)?;
        f(output)
    }

    /// Apply module and map the output with a function
    ///
    /// Similar to `and_then` but for non-failable transformations.
    ///
    /// # Arguments
    /// * `input` - Input tensor
    /// * `f` - Mapping function
    ///
    /// # Returns
    /// * `Result<Tensor>` - Mapped output
    fn map<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
    where
        F: FnOnce(Tensor) -> Tensor,
    {
        let output = self.forward(input)?;
        Ok(f(output))
    }

    /// Forward pass with input transformation
    ///
    /// Apply a transformation to the input before forwarding.
    ///
    /// # Arguments
    /// * `input` - Input tensor
    /// * `f` - Input transformation function
    ///
    /// # Returns
    /// * `Result<Tensor>` - Module output
    fn with_input<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
    where
        F: FnOnce(&Tensor) -> Result<Tensor>,
    {
        let transformed = f(input)?;
        self.forward(&transformed)
    }

    // === Inspection and Debugging Methods ===

    /// Get human-readable summary of the module
    ///
    /// # Returns
    /// * `String` - Formatted module summary
    fn summary(&self) -> String {
        let info = self.module_info();
        format!(
            "Module: {}\n\
             Training: {}\n\
             Parameters: {} ({} trainable)\n\
             Memory: {:.2} MB\n\
             Children: {}",
            info.name,
            info.training,
            info.parameter_count,
            info.trainable_parameter_count,
            info.memory_usage_bytes as f64 / (1024.0 * 1024.0),
            info.children_count
        )
    }

    /// Print module summary to stdout
    fn print_summary(&self) {
        println!("{}", self.summary());
    }

    /// Get parameter statistics
    ///
    /// # Returns
    /// * `ParameterStats` - Statistical information about parameters
    fn parameter_stats(&self) -> ParameterStats {
        let params = self.all_parameters();
        let mut total_params = 0;
        let mut trainable_params = 0;
        let mut frozen_params = 0;
        let mut total_memory = 0;

        for param in params.values() {
            let numel = param.numel().unwrap_or(0);
            total_params += numel;
            total_memory += numel * 4; // Assume f32

            if param.requires_grad() {
                trainable_params += numel;
            } else {
                frozen_params += numel;
            }
        }

        ParameterStats {
            total_parameters: total_params,
            trainable_parameters: trainable_params,
            frozen_parameters: frozen_params,
            total_memory_bytes: total_memory,
            parameter_count: params.len(),
        }
    }

    /// Check if module has NaN or Inf in parameters
    ///
    /// # Returns
    /// * `bool` - true if all parameters are finite
    fn has_finite_parameters(&self) -> bool {
        self.all_parameters()
            .values()
            .all(|p| p.is_finite().unwrap_or(false))
    }

    /// Get list of parameter names
    ///
    /// # Returns
    /// * `Vec<String>` - Sorted list of parameter names
    fn parameter_names(&self) -> Vec<String> {
        let mut names: Vec<String> = self.all_named_parameters().keys().cloned().collect();
        names.sort();
        names
    }

    /// Get parameter by name
    ///
    /// # Arguments
    /// * `name` - Parameter name
    ///
    /// # Returns
    /// * `Option<Parameter>` - Parameter if found
    fn get_parameter(&self, name: &str) -> Option<crate::Parameter> {
        self.all_named_parameters().get(name).cloned()
    }

    // === Training Utilities ===

    /// Freeze parameters whose name contains `pattern` (sets `requires_grad = false`)
    ///
    /// Pass an empty pattern (`""`) to freeze every parameter, since every name
    /// contains the empty string.
    ///
    /// # Arguments
    /// * `pattern` - Substring to match against parameter names
    ///
    /// # Returns
    /// * `usize` - Number of parameters frozen
    ///
    /// # Note
    /// A parameter's `requires_grad` flag is shared across clones (see
    /// [`crate::Parameter`]), so updating the by-value clones handed out by
    /// [`Module::all_named_parameters`] is observable on the module's own
    /// parameters and on any subsequently fetched copies.
    fn freeze_matching(&mut self, pattern: &str) -> usize {
        let mut count = 0;
        for (name, param) in self.all_named_parameters() {
            if name.contains(pattern) {
                param.set_requires_grad(false);
                count += 1;
            }
        }
        count
    }

    /// Unfreeze parameters whose name contains `pattern` (sets `requires_grad = true`)
    ///
    /// Pass an empty pattern (`""`) to unfreeze every parameter, since every name
    /// contains the empty string.
    ///
    /// # Arguments
    /// * `pattern` - Substring to match against parameter names
    ///
    /// # Returns
    /// * `usize` - Number of parameters unfrozen
    ///
    /// # Note
    /// See [`ModuleExt::freeze_matching`] for why mutating the by-value parameter
    /// clones is observable on the owning module.
    fn unfreeze_matching(&mut self, pattern: &str) -> usize {
        let mut count = 0;
        for (name, param) in self.all_named_parameters() {
            if name.contains(pattern) {
                param.set_requires_grad(true);
                count += 1;
            }
        }
        count
    }

    /// Get list of frozen parameters
    ///
    /// # Returns
    /// * `Vec<String>` - Names of frozen parameters
    fn frozen_parameters(&self) -> Vec<String> {
        self.all_named_parameters()
            .into_iter()
            .filter(|(_, p)| !p.requires_grad())
            .map(|(name, _)| name)
            .collect()
    }

    /// Get list of trainable parameters
    ///
    /// # Returns
    /// * `Vec<String>` - Names of trainable parameters
    fn trainable_parameters(&self) -> Vec<String> {
        self.all_named_parameters()
            .into_iter()
            .filter(|(_, p)| p.requires_grad())
            .map(|(name, _)| name)
            .collect()
    }

    // === Advanced Operations ===

    /// Clone module parameters into a new state dict
    ///
    /// # Returns
    /// * `HashMap<String, Tensor>` - Cloned state dictionary
    fn clone_state_dict(&self) -> HashMap<String, Tensor> {
        self.state_dict()
    }

    /// Apply a function to all parameters
    ///
    /// # Arguments
    /// * `f` - Function to apply to each parameter
    fn apply_to_parameters<F>(&self, mut f: F)
    where
        F: FnMut(&str, &crate::Parameter),
    {
        for (name, param) in self.all_named_parameters() {
            f(&name, &param);
        }
    }

    /// Count parameters by layer type
    ///
    /// # Returns
    /// * `HashMap<String, usize>` - Parameter count per layer type
    fn parameters_by_type(&self) -> HashMap<String, usize> {
        let mut counts = HashMap::new();

        for (name, param) in self.all_named_parameters() {
            // Extract layer type from name (first component)
            let layer_type = name.split('.').next().unwrap_or("unknown").to_string();

            let numel = param.numel().unwrap_or(0);
            *counts.entry(layer_type).or_insert(0) += numel;
        }

        counts
    }

    /// Validate module configuration
    ///
    /// Performs comprehensive validation of module state.
    ///
    /// # Returns
    /// * `Result<ValidationReport>` - Validation results
    fn validate(&self) -> Result<ValidationReport> {
        let mut report = ValidationReport::default();

        // Check for parameters
        if !self.has_parameters() {
            report.warnings.push("Module has no parameters".to_string());
        }

        // Check for finite parameters
        if !self.has_finite_parameters() {
            report
                .errors
                .push("Module has non-finite parameters (NaN or Inf)".to_string());
        }

        // Check memory usage
        let memory_mb = self.memory_usage_mb();
        if memory_mb > 1024.0 {
            report
                .warnings
                .push(format!("Large memory usage: {:.2} GB", memory_mb / 1024.0));
        }

        // Check parameter count
        let param_count = self.num_parameters();
        if param_count > 100_000_000 {
            report
                .warnings
                .push(format!("Very large model: {} parameters", param_count));
        }

        report.is_valid = report.errors.is_empty();
        Ok(report)
    }

    /// Get the device the module's parameters reside on
    ///
    /// # Returns
    /// * `Option<DeviceType>` - The device of the module's parameters, or `None`
    ///   only when the module genuinely has no parameters (recursively).
    ///
    /// # Note
    /// The device is read from the actual underlying parameter tensors. If a
    /// module's parameters span multiple devices, the device of the first
    /// parameter encountered is reported.
    fn device(&self) -> Option<DeviceType> {
        self.all_parameters().values().next().map(|p| p.device())
    }

    /// Check if all parameters are on CPU
    ///
    /// # Returns
    /// * `bool` - true if all parameters on CPU
    fn is_cpu(&self) -> bool {
        self.device() == Some(DeviceType::Cpu)
    }

    /// Check if all parameters are on CUDA device
    ///
    /// # Returns
    /// * `bool` - true if all parameters on CUDA
    fn is_cuda(&self) -> bool {
        matches!(self.device(), Some(DeviceType::Cuda(_)))
    }
}

// Automatically implement ModuleExt for all types that implement Module
impl<T: Module + ?Sized> ModuleExt for T {}

// === Supporting Types ===

/// Parameter statistics for a module
#[derive(Debug, Clone)]
pub struct ParameterStats {
    /// Total number of parameter elements
    pub total_parameters: usize,
    /// Number of trainable parameter elements
    pub trainable_parameters: usize,
    /// Number of frozen parameter elements
    pub frozen_parameters: usize,
    /// Total memory usage in bytes
    pub total_memory_bytes: usize,
    /// Number of distinct parameters
    pub parameter_count: usize,
}

impl ParameterStats {
    /// Get memory usage in megabytes
    pub fn memory_mb(&self) -> f64 {
        self.total_memory_bytes as f64 / (1024.0 * 1024.0)
    }

    /// Get memory usage in gigabytes
    pub fn memory_gb(&self) -> f64 {
        self.memory_mb() / 1024.0
    }

    /// Get percentage of parameters that are trainable
    pub fn trainable_percentage(&self) -> f64 {
        if self.total_parameters == 0 {
            0.0
        } else {
            (self.trainable_parameters as f64 / self.total_parameters as f64) * 100.0
        }
    }
}

/// Validation report for a module
#[derive(Debug, Clone, Default)]
pub struct ValidationReport {
    /// Whether the module is valid
    pub is_valid: bool,
    /// List of errors found
    pub errors: Vec<String>,
    /// List of warnings
    pub warnings: Vec<String>,
}

impl ValidationReport {
    /// Check if validation passed without errors
    pub fn passed(&self) -> bool {
        self.is_valid && self.errors.is_empty()
    }

    /// Get total number of issues (errors + warnings)
    pub fn issue_count(&self) -> usize {
        self.errors.len() + self.warnings.len()
    }

    /// Format as human-readable string
    pub fn format(&self) -> String {
        let mut result = String::new();

        result.push_str(&format!(
            "Validation: {}\n",
            if self.is_valid { "PASSED" } else { "FAILED" }
        ));

        if !self.errors.is_empty() {
            result.push_str("\nErrors:\n");
            for error in &self.errors {
                result.push_str(&format!("  - {}\n", error));
            }
        }

        if !self.warnings.is_empty() {
            result.push_str("\nWarnings:\n");
            for warning in &self.warnings {
                result.push_str(&format!("  - {}\n", warning));
            }
        }

        result
    }
}

#[cfg(test)]
mod tests {
    use super::ModuleExt;
    use crate::layers::Linear;
    use crate::Module;
    use torsh_core::device::DeviceType;
    use torsh_core::error::Result;
    use torsh_tensor::Tensor;

    /// A module that genuinely owns no parameters, used to verify that `device()`
    /// honestly reports `None` instead of a fabricated default.
    struct EmptyModule;

    impl Module for EmptyModule {
        fn forward(&self, input: &Tensor) -> Result<Tensor> {
            Ok(input.clone())
        }
    }

    #[test]
    fn test_freeze_unfreeze_sets_requires_grad() {
        let mut linear = Linear::new(4, 3, true);

        // A Linear(.., bias = true) layer exposes exactly two parameters.
        let param_count = linear.all_parameters().len();
        assert_eq!(
            param_count, 2,
            "Linear with bias should expose weight + bias"
        );

        // Freshly created parameters must require gradients.
        assert!(
            linear.all_parameters().values().all(|p| p.requires_grad()),
            "new Linear parameters should require grad"
        );

        // Freeze every parameter (empty pattern matches all names).
        let frozen = linear.freeze_matching("");
        assert_eq!(
            frozen, param_count,
            "freeze_matching(\"\") must touch every parameter"
        );

        // The freeze must be observable on a *fresh* read of the parameters,
        // proving the flag really mutated shared state (not a throwaway clone).
        assert!(
            linear.all_parameters().values().all(|p| !p.requires_grad()),
            "after freeze no parameter may require grad"
        );
        assert_eq!(
            linear.frozen_parameters().len(),
            param_count,
            "all parameters should be reported as frozen"
        );
        assert!(
            linear.trainable_parameters().is_empty(),
            "no parameter should be trainable after freeze"
        );

        // Unfreeze every parameter and verify it sticks.
        let unfrozen = linear.unfreeze_matching("");
        assert_eq!(unfrozen, param_count);
        assert!(
            linear.all_parameters().values().all(|p| p.requires_grad()),
            "after unfreeze every parameter must require grad"
        );
        assert!(linear.frozen_parameters().is_empty());
    }

    #[test]
    fn test_freeze_matching_only_affects_matching_names() {
        let mut linear = Linear::new(4, 3, true);

        // Freeze only the bias; the weight must stay trainable.
        let frozen = linear.freeze_matching("bias");
        assert_eq!(frozen, 1, "only the bias parameter should match");

        let params = linear.all_named_parameters();
        assert!(
            !params["bias"].requires_grad(),
            "bias must be frozen after freeze_matching(\"bias\")"
        );
        assert!(
            params["weight"].requires_grad(),
            "weight must remain trainable"
        );
    }

    #[test]
    fn test_device_reports_real_parameter_device() {
        let linear = Linear::new(8, 5, true);

        // Parameters are created on CPU, so the module device must be CPU.
        assert_eq!(linear.device(), Some(DeviceType::Cpu));
        assert!(linear.is_cpu());
        assert!(!linear.is_cuda());
    }

    #[test]
    fn test_device_is_none_without_parameters() {
        let empty = EmptyModule;
        // A parameter-less module must honestly report None, never a fake device.
        assert_eq!(empty.device(), None);
        assert!(!empty.is_cpu());
        assert!(!empty.is_cuda());
    }
}