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 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.

Source

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

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.

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 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.

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".

Implementors§

Source§

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