Skip to main content

ModuleMapper

burn_core::module

Trait ModuleMapper 

Source 
pub trait ModuleMapper {
    // Provided methods
    fn enter_module(&mut self, name: &str, container_type: &str) { ... }
    fn exit_module(&mut self, name: &str, container_type: &str) { ... }
    fn map_float<const D: usize>(
        &mut self,
        param: Param<Tensor<D>>,
    ) -> Param<Tensor<D>> { ... }
    fn map_int<const D: usize>(
        &mut self,
        param: Param<Tensor<D, Int>>,
    ) -> Param<Tensor<D, Int>> { ... }
    fn map_bool<const D: usize>(
        &mut self,
        param: Param<Tensor<D, Bool>>,
    ) -> Param<Tensor<D, Bool>> { ... }
    fn map_flag(&mut self, flag: Param<Flag>) -> Param<Flag> { ... }
}
Expand description

Module mapper trait for transforming module parameters.

Provided Methods§

Source

fn enter_module(&mut self, name: &str, container_type: &str)

Called when entering a submodule.

§Parameters
  • name: The name of the submodule being entered
  • container_type: The type of the container with format:
    • For user-defined structs: “Struct:TypeName” (e.g., “Struct:Linear”)
    • For user-defined enums: “Enum:TypeName” (e.g., “Enum:MyEnum”)
    • For Vec containers: “Vec” (name is the index)
    • For Tuple containers: “Tuple” (name is the index)
    • For Array containers: “Array” (name is the index)

Note: Option containers do not call enter_module/exit_module to preserve the field name in the path (e.g., “bias” instead of “bias.Some”)

Source

fn exit_module(&mut self, name: &str, container_type: &str)

Called when exiting a submodule.

§Parameters
  • name: The name of the submodule being exited
  • container_type: The type of the container with format:
    • For user-defined structs: “Struct:TypeName” (e.g., “Struct:Linear”)
    • For user-defined enums: “Enum:TypeName” (e.g., “Enum:MyEnum”)
    • For Vec containers: “Vec” (name is the index)
    • For Tuple containers: “Tuple” (name is the index)
    • For Array containers: “Array” (name is the index)

Note: Option containers do not call enter_module/exit_module to preserve the field name in the path (e.g., “bias” instead of “bias.Some”)

Source

fn map_float<const D: usize>( &mut self, param: Param<Tensor<D>>, ) -> Param<Tensor<D>>

Map a float parameter in the module.

§Parameters
  • param: The float parameter to transform
§Returns

The transformed parameter

Source

fn map_int<const D: usize>( &mut self, param: Param<Tensor<D, Int>>, ) -> Param<Tensor<D, Int>>

Map an int parameter in the module.

§Parameters
  • param: The integer parameter to transform
§Returns

The transformed parameter

Source

fn map_bool<const D: usize>( &mut self, param: Param<Tensor<D, Bool>>, ) -> Param<Tensor<D, Bool>>

Map a bool parameter in the module.

§Parameters
  • param: The boolean parameter to transform
§Returns

The transformed parameter

Source

fn map_flag(&mut self, flag: Param<Flag>) -> Param<Flag>

Map a Param<Flag> in the module.

The default is the identity, so a mapper with no opinion about control state (for example, a record, a quantizer or a module sharder) leaves it untouched. Device moves are handled directly by Module::to_device and do not go through a ModuleMapper.

§Parameters
  • flag: The flag to map
§Returns

The mapped flag

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§