Skip to main content

ParamReshaper

Trait ParamReshaper 

Source
pub trait ParamReshaper<B: Backend>: Send + Sync {
    type Module: Module<B>;

    // Required methods
    fn num_params(&self) -> usize;
    fn flatten(&self, module: &Self::Module, device: &B::Device) -> Tensor<B, 1>;
    fn unflatten(&self, flat: Tensor<B, 1>) -> Self::Module;
}
Expand description

Bridges a Burn Module and a flat Tensor<B, 1> parameter vector.

Lives entirely in rlevo-evolution and depends only on burn — no rlevo-core coupling.

§Invariants

  • flatten and unflatten must visit float leaves in the same deterministic order, so that unflatten(flatten(m)) reconstructs m leaf-for-leaf.
  • num_params equals the total element count produced by flatten.

Implementors are Send + Sync so a single reshaper can be shared across parallel fitness evaluations.

Required Associated Types§

Source

type Module: Module<B>

The Burn module type this reshaper flattens and reconstructs.

Required Methods§

Source

fn num_params(&self) -> usize

Total number of trainable float parameters (the flat-vector length).

Source

fn flatten(&self, module: &Self::Module, device: &B::Device) -> Tensor<B, 1>

Flatten all float Param leaves of module into a 1-D tensor.

The returned tensor is moved onto device and has length num_params. Leaf visitation order is deterministic and matches unflatten.

§Panics

Panics if module has no float leaves (the underlying tensor concatenation requires at least one part).

Source

fn unflatten(&self, flat: Tensor<B, 1>) -> Self::Module

Clone the template module and replace its float leaves with slices of flat, in the same order as flatten.

§Aliasing

The returned module’s leaves are views into flat’s storage (Burn Tensor::clone is a refcount bump; slice/reshape are view ops). This is allocation-free and safe for the forward-only scoring this bridge exists for. Do not mutate flat or the returned module’s weights in place while the other is live — they share backing storage. The output module inherits flat’s device.

§Panics

Panics if flat.dims()[0] != self.num_params().

Dyn Compatibility§

This trait is dyn compatible.

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

Implementors§

Source§

impl<B, M> ParamReshaper<B> for ModuleReshaper<B, M>
where B: Backend, M: Module<B> + Sync,