Struct Flatten 

pub struct Flatten { /* private fields */ }

Flatten layer: collapses dimensions [start_dim..=end_dim] into one.

Default (start_dim=1): [N, C, H, W] → [N, C*H*W].

Examples

let flatten = Flatten::new(1); // flatten from dim 1 onward
let x: [2, 8, 4, 4] tensor
let y = flatten.forward(&x)?; // [2, 128]

Implementations

impl Flatten

pub fn new(start_dim: usize) -> Flatten

Create a Flatten that collapses from start_dim through the last dim.

• start_dim = 1: flatten everything except batch → [N, ...] → [N, flat]
• start_dim = 0: flatten everything → [total]

pub fn default_flat() -> Flatten

Default flatten: from dim 1 onward (preserves batch dimension).

Trait Implementations

impl Default for Flatten

fn default() -> Flatten

Returns the “default value” for a type.

impl<B> Module<B> for Flatten

where B: Backend,

fn forward(&self, x: &Tensor<B>) -> Result<Tensor<B>, Error>

Compute the output tensor from the input tensor. This defines the layer’s computation (forward pass).

fn parameters(&self) -> Vec<Tensor<B>>

Return all trainable parameters of this module. The optimizer uses these to update weights during training.

fn set_training(&self, _training: bool)

Set training or evaluation mode.

fn is_training(&self) -> bool

Whether the module is in training mode (default: true).

fn train(&self)

Convenience: set training mode.

fn eval(&self)

Convenience: set evaluation mode.

fn num_parameters(&self) -> usize

Total number of scalar parameters in this module.

fn trainable_params_count(&self) -> usize

Number of trainable (variable) parameters.

fn frozen_parameters(&self) -> Vec<Tensor<B>>

Freeze all parameters: returns new parameter tensors with is_variable = false, preventing gradient accumulation.

fn named_parameters(&self) -> Vec<(String, Tensor<B>)>

Return all trainable parameters with human-readable names.

fn state_dict(&self) -> Vec<(String, Tensor<B>)>

Returns a state_dict-style map of parameter name → tensor.