memorax.networks.blocks

memorax.networks.blocks#

Building blocks for constructing network architectures.

Feed-Forward#

FFN - Feed-forward network block with expansion.

GLU - Gated Linear Unit variant of FFN.

Projection - Single linear projection layer.

Normalization#

PreNorm - Pre-normalization wrapper.

PostNorm - Post-normalization wrapper.

Residual#

Residual - Residual connection wrapper.

GatedResidual - Gated residual connection with learned gate.

Composition#

Stack - Stacks multiple blocks sequentially.

SegmentRecurrence - Fixed-length cross-segment memory buffer.

Mixture of Experts#

MoE - Mixture of Experts layer.

TopKRouter - Top-K routing for MoE.