Struct Layer

pub struct Layer<M>where
    M: Module,{
    pub norm: RmsNorm,
    pub block: M,
    pub norm2: Option<RmsNorm>,
    pub mlp: Option<GatedMlp>,
    pub class_latents: Vec<ClassLatent>,
    pub class_latents_emb: Option<Param<Tensor<2>>>,
}

Expand description

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.

With Self::mlp set the layer additionally runs a second Pre-LN sub-block, a SwiGLU feed-forward (see GatedMlp). It has a residual of its own, inside the layer, which is the reason the methods below return the layer’s total delta rather than the mixer output:

  h₁ = M(norm(x))                     the mixer sub-block
  h₂ = mlp(norm2(x + h₁))             the feed-forward sub-block
  return h₁ + h₂                      so that Layers' `x + delta` is
                                      (x + h₁) + h₂ — both residuals

Folding it this way keeps Layers the single owner of the outer residual (and of the ignore_first/last_residual ablations, which therefore govern only that outer add — the feed-forward’s inner residual is intrinsic to the sub-block and always applies). Without an mlp the delta is just h₁ and nothing changes for a block family that carries no feed-forward.

May carry its own ClassLatents, placed from a [ClassCursor]: step splices them around the token it is given, while in forward the caller splices them first (via Self::insert_latents) so the residual it adds sees the same lengthened sequence; Self::prime steps the ones waiting for the next token without that token. They are independent of any class latents on the enclosing Layers.

Fields§

§norm: RmsNorm

Pre-norm applied before the inner block.

§block: M

The inner mixer block.

§norm2: Option<RmsNorm>

Pre-norm of the feed-forward sub-block. Some exactly when Self::mlp is (norm2 in the reference checkpoints).

§mlp: Option<GatedMlp>

Optional SwiGLU feed-forward sub-block run after the mixer, with its own residual. None ⇒ the layer is mixer-only.

§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> Layer<M>
where M: Block,

pub fn insert_latents( &self, x: Tensor<3>, class: Option<&mut ClassCursor>, ) -> Tensor<3>

Splice this layer’s class latents into the chunk x (no-op when there are none), advancing class past it.

Public so a caller driving a bare Layer can lengthen the sequence itself (and add the matching residual) before calling Self::forward. None cursors ⇒ this chunk is the whole sequence. Layers splices its layers’ latents itself, since under MultiGate residuals the same rows must also enter the carried streams.

pub fn forward( &self, x: Tensor<3>, cache: Option<<M as Block>::Cache>, options: <M as Block>::Options, ) -> (Tensor<3>, <M as Block>::Cache)

Full-sequence Pre-LN block without the outer residual: the layer’s total delta M(RMSNorm(x)), plus the feed-forward sub-block’s own contribution when Self::mlp is set (see the type docs).

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

pub fn step( &self, x: Tensor<2>, cache: Option<<M as Block>::Cache>, class: Option<&mut ClassCursor>, ) -> (Tensor<2>, <M as Block>::Cache)

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

class is this layer’s own class-latent cursor. With Some, every latent whose position falls on this token is stepped around it — before it (Start/Middle/Custom, which precede a token) or after it (End, which closes the sequence) — each a step of its own. What comes back is the last token the step emitted (see ClassCursors): the user token, unless an End latent follows it, that latent being then the sequence’s true last token. With None no class latents are injected — and Middle/End latents panic (their positions need the full sequence length). The residual is the caller’s responsibility.

pub fn prime( &self, batch: usize, cache: Option<<M as Block>::Cache>, class: Option<&mut ClassCursor>, ) -> (Option<(Tensor<2>, Tensor<2>)>, Option<<M as Block>::Cache>)

Step the class latents this layer has waiting for its next token — with no token of its own, so nothing but class data is consumed.

This is Self::step’s opening half on its own (see ClassCursors): the latents that would have preceded the next token are stepped now, in the same order, so a prime followed by a step runs exactly the sequence that step alone would have. End latents are never primed — closing the sequence, they belong to the step carrying its last token.

Returns the last latent stepped, as the pair (delta, latent) — this layer’s own embedding row alongside the delta it produced, since the caller has no other way to complete the residual (delta + latent, as it does with the token it hands to Self::step). None ⇒ nothing was waiting, and the cache comes back exactly as it went in (None included: a layer that stepped nothing has the state it already had).

pub fn step_one( &self, x: Tensor<2>, cache: Option<<M as Block>::Cache>, ) -> (Tensor<2>, <M as Block>::Cache)

The actual one-token work: no class injection, no outer residual.

Layers’s cascade uses it to place this layer’s class latents from the stack-wide ClassCursors itself, bypassing Self::step’s cursorless guard (that guard rejects Middle/End, which the cascade has already resolved). It is public because an external container that owns the residual — one threading its own state between layers rather than a per-layer cache — needs exactly this: the layer’s delta and the cache it produced, with nothing added.

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 Block::block_step_infinite). Cursorless: class latents are not injected (Middle/End latents panic, as in a None-cursor step). The feed-forward sub-block is point-wise, so it composes with the limit: once the mixer output settles, x + h₁ is constant and so is h₂.