Skip to main content

ModuleExt

torsh_nn::core::module_ext

Trait ModuleExt 

Source 
pub trait ModuleExt: Module {

Show 20 methods    // Provided methods
    fn and_then<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
       where F: FnOnce(Tensor) -> Result<Tensor> { ... }
    fn map<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
       where F: FnOnce(Tensor) -> Tensor { ... }
    fn with_input<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
       where F: FnOnce(&Tensor) -> Result<Tensor> { ... }
    fn summary(&self) -> String { ... }
    fn print_summary(&self) { ... }
    fn parameter_stats(&self) -> ParameterStats { ... }
    fn has_finite_parameters(&self) -> bool { ... }
    fn parameter_names(&self) -> Vec<String> { ... }
    fn get_parameter(&self, name: &str) -> Option<Parameter> { ... }
    fn freeze_matching(&mut self, pattern: &str) -> usize { ... }
    fn unfreeze_matching(&mut self, pattern: &str) -> usize { ... }
    fn frozen_parameters(&self) -> Vec<String> { ... }
    fn trainable_parameters(&self) -> Vec<String> { ... }
    fn clone_state_dict(&self) -> HashMap<String, Tensor> { ... }
    fn apply_to_parameters<F>(&self, f: F)
       where F: FnMut(&str, &Parameter) { ... }
    fn parameters_by_type(&self) -> HashMap<String, usize> { ... }
    fn validate(&self) -> Result<ValidationReport> { ... }
    fn device(&self) -> Option<DeviceType> { ... }
    fn is_cpu(&self) -> bool { ... }
    fn is_cuda(&self) -> bool { ... }

}
Expand description

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

Provided Methods§

Source

fn and_then<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
where F: FnOnce(Tensor) -> Result<Tensor>,

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
ⓘ
let output = layer.and_then(&input, |x| x.relu())?;
Source

fn map<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
where F: FnOnce(Tensor) -> Tensor,

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
Source

fn with_input<F>(&self, input: &Tensor, f: F) -> Result<Tensor>
where F: FnOnce(&Tensor) -> Result<Tensor>,

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
Source

fn summary(&self) -> String

Get human-readable summary of the module

§Returns
  • String - Formatted module summary
Source

fn print_summary(&self)

Print module summary to stdout

Source

fn parameter_stats(&self) -> ParameterStats

Get parameter statistics

§Returns
  • ParameterStats - Statistical information about parameters
Source

fn has_finite_parameters(&self) -> bool

Check if module has NaN or Inf in parameters

§Returns
  • bool - true if all parameters are finite
Source

fn parameter_names(&self) -> Vec<String>

Get list of parameter names

§Returns
  • Vec<String> - Sorted list of parameter names
Source

fn get_parameter(&self, name: &str) -> Option<Parameter>

Get parameter by name

§Arguments
  • name - Parameter name
§Returns
  • Option<Parameter> - Parameter if found
Source

fn freeze_matching(&mut self, pattern: &str) -> usize

Freeze specific parameters by name pattern

§Arguments
  • pattern - String pattern to match parameter names
§Returns
  • usize - Number of parameters frozen
§Note

This method currently returns a count but doesn’t actually freeze parameters because Parameter’s requires_grad is immutable. This is a placeholder for future implementation when mutable parameter access is available.

Source

fn unfreeze_matching(&mut self, pattern: &str) -> usize

Unfreeze specific parameters by name pattern

§Arguments
  • pattern - String pattern to match parameter names
§Returns
  • usize - Number of parameters unfrozen
§Note

This method currently returns a count but doesn’t actually unfreeze parameters because Parameter’s requires_grad is immutable. This is a placeholder for future implementation when mutable parameter access is available.

Source

fn frozen_parameters(&self) -> Vec<String>

Get list of frozen parameters

§Returns
  • Vec<String> - Names of frozen parameters
Source

fn trainable_parameters(&self) -> Vec<String>

Get list of trainable parameters

§Returns
  • Vec<String> - Names of trainable parameters
Source

fn clone_state_dict(&self) -> HashMap<String, Tensor>

Clone module parameters into a new state dict

§Returns
  • HashMap<String, Tensor> - Cloned state dictionary
Source

fn apply_to_parameters<F>(&self, f: F)
where F: FnMut(&str, &Parameter),

Apply a function to all parameters

§Arguments
  • f - Function to apply to each parameter
Source

fn parameters_by_type(&self) -> HashMap<String, usize>

Count parameters by layer type

§Returns
  • HashMap<String, usize> - Parameter count per layer type
Source

fn validate(&self) -> Result<ValidationReport>

Validate module configuration

Performs comprehensive validation of module state.

§Returns
  • Result<ValidationReport> - Validation results
Source

fn device(&self) -> Option<DeviceType>

Get device of parameters (if consistent)

§Returns
  • Option<DeviceType> - Device if all parameters are on same device
§Note

Currently returns None as Parameter doesn’t expose device information. This is a placeholder for future implementation.

Source

fn is_cpu(&self) -> bool

Check if all parameters are on CPU

§Returns
  • bool - true if all parameters on CPU
Source

fn is_cuda(&self) -> bool

Check if all parameters are on CUDA device

§Returns
  • bool - true if all parameters on CUDA

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety", so this trait is not object safe.

Implementors§

Source§

impl<T: Module + ?Sized> ModuleExt for T