Struct Layer 

pub struct Layer<M: Module> {
    pub norm: RmsNorm,
    pub mamba_block: M,
    pub class_latents: Vec<ClassLatent>,
    pub class_latents_emb: Option<Param<Tensor<2>>>,
}

A single Pre-LN block wrapper computing M(RMSNorm(x)) — the residual is not applied here. The enclosing Layers owns that decision (add the input back, suppress it on the first/last layer, or thread it through Multi-Gate streams), so no input clone / zero-add is wasted when no residual is wanted.

May carry its own ClassLatents. In step they are spliced via the index cursor; in forward the caller splices them first (via Self::insert_latents) so the residual it adds sees the same lengthened sequence. They are independent of any class latents on the enclosing Layers.

Fields

norm: RmsNorm

Pre-norm applied before the inner block.

mamba_block: M

The inner Mamba-x SSM block.

class_latents: Vec<ClassLatent>

Positions of this layer’s class latents (empty ⇒ none).

class_latents_emb: Option<Param<Tensor<2>>>

The class-latent embeddings, [num_class_latents, d_model] (None ⇒ none).

Implementations

impl<M: MambaBlock> Layer<M>

pub(crate) fn insert_latents(&self, x: Tensor<3>) -> Tensor<3>

Splice this layer’s class latents into x (no-op when there are none). Public to the crate so Layers can lengthen the sequence itself (and add the matching residual) before calling Self::forward.

pub fn forward( &self, x: Tensor<3>, cache: Option<M::Cache>, ssd_path: M::SsdPath, ) -> (Tensor<3>, M::Cache)

Full-sequence Pre-LN block without the residual: M(RMSNorm(x)).

The caller owns any class-latent insertion (Self::insert_latents) and the residual.

pub fn step( &self, x: Tensor<2>, cache: Option<M::Cache>, index: Option<&mut usize>, ) -> (Tensor<2>, M::Cache)

Single-token Pre-LN block step without the residual.

index is the running cursor into this layer’s output sequence. With Some, whenever it lands on one of this layer’s class-latent positions those latents are stepped first (each advancing index, recursing with None); only the user token’s output and cache are returned. With None no class latents are injected — and Middle/End latents panic (their positions need the full sequence; use forward). The residual is the caller’s responsibility.

pub fn step_infinite(&self, x: Tensor<2>) -> Tensor<2>

Stationary fixed point of the Pre-LN block under a constant token, without the residual: the step counterpart of infinitely many identical tokens (closed form, no cache — see MambaBlock::block_step_infinite). Cursorless: class latents are not injected (Middle/End latents panic, as in a None-cursor step).

pub fn step_n_approx( &self, x: Tensor<2>, n: usize, cache: Option<M::Cache>, ) -> (Tensor<2>, M::Cache)

Closed-form jump equivalent to n cursorless Self::step calls on the same constant token, without the residual (see MambaBlock::block_step_n_approx).

Trait Implementations

impl<M> AutodiffModule for Layer<M>

where M: AutodiffModule + ModuleDisplay + Module,

fn valid(&self) -> Self

Returns the same module, but on the inner backend without auto-differentiation.

fn from_inner(module: Self) -> Self

Wraps an inner module back into an auto-diff module.

impl<M> Clone for Layer<M>

where M: Module + ModuleDisplay + Module,

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<M: Debug + Module> Debug for Layer<M>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<M> Display for Layer<M>

where M: Module + ModuleDisplay + Module,

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<M> Module for Layer<M>

where M: Module + ModuleDisplay + Module,

type Record = LayerRecord<M>

Type to save and load the module.

fn load_record(self, record: Self::Record) -> Self

Load the module state from a record.

fn into_record(self) -> Self::Record

Convert the module into a record containing the state.

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules.

fn visit<Visitor: ModuleVisitor>(&self, visitor: &mut Visitor)

Visit each tensor parameter in the module with a visitor.

fn map<Mapper: ModuleMapper>(self, mapper: &mut Mapper) -> Self

Map each tensor parameter in the module with a mapper.

fn collect_devices(&self, devices: Devices) -> Devices

Return all the devices found in the underneath module tree added to the given vector without duplicates.

fn to_device(self, device: &Device) -> Self

Move the module and all of its sub-modules to the given device.

fn fork(self, device: &Device) -> Self

Fork the module and all of its sub-modules to the given device.

fn devices(&self) -> Vec<Device>

Return all the devices found in the underneath module tree without duplicates.

fn no_grad(self) -> Self

Each tensor in the module tree will not require grad.

fn train(self) -> Self

where Self: AutodiffModule,

Move the module and all of its sub-modules to the autodiff backend.

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

Quantize the weights of the module.

impl<M> ModuleDisplay for Layer<M>

where M: Module + ModuleDisplay + Module,

fn format(&self, passed_settings: DisplaySettings) -> String

Formats the module with provided display settings.

fn custom_settings(&self) -> Option<DisplaySettings>

Custom display settings for the module.

fn custom_content(&self, _content: Content) -> Option<Content>

Custom attributes for the module.

impl<M> ModuleDisplayDefault for Layer<M>

where M: Module + ModuleDisplay + Module,

fn content(&self, content: Content) -> Option<Content>

Attributes of the module used for display purposes.

fn num_params(&self) -> usize

Gets the number of the parameters of the module.