Struct BidiLayerPair 

pub struct BidiLayerPair<M: Module> {
    pub straight_norm: RmsNorm,
    pub reverse_norm: RmsNorm,
    pub straight_block: M,
    pub reverse_block: M,
    pub output_merge: OutputMerge,
    pub class_latents: Vec<ClassLatent>,
    pub class_latents_emb: Option<Param<Tensor<2>>>,
}

A single bidirectional pair: a straight (→) and a reversed (←) Pre-LN block whose outputs are merged. The residual is not applied here — the enclosing BidiLayers adds it (or suppresses it on the first/last pair), mirroring the Layer / Layers split.

Fields

straight_norm: RmsNorm

Pre-norm for the straight pass.

reverse_norm: RmsNorm

Pre-norm for the reversed pass.

straight_block: M

The block run left-to-right.

reverse_block: M

The block run right-to-left (over the flipped sequence).

output_merge: OutputMerge

Merge strategy combining the two directions.

class_latents: Vec<ClassLatent>

Positions of this pair’s class latents, spliced in before either direction runs (both directions, and the residual, see the lengthened sequence). Empty ⇒ none.

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

This pair’s class-latent embeddings, [num_class_latents, d_model].

Implementations

impl<M: MambaBlock> BidiLayerPair<M>

where M::SsdPath: Clone,

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

Splice this bidi-layer-pair’s class latents into x (no-op when there are none).

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

[batch, sequence, d_model] → [batch, sequence, d_model], plus the two updated direction caches. (sequence grows by the class-latent count.) Returns the merged directions without the residual — the enclosing BidiLayers adds it.

Trait Implementations

impl<M> AutodiffModule for BidiLayerPair<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 BidiLayerPair<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 BidiLayerPair<M>

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

Formats the value using the given formatter.

impl<M> Display for BidiLayerPair<M>

where M: Module + ModuleDisplay + Module,

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

Formats the value using the given formatter.

impl<M> Module for BidiLayerPair<M>

where M: Module + ModuleDisplay + Module,

type Record = BidiLayerPairRecord<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 BidiLayerPair<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 BidiLayerPair<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.