memorax.networks.layers

Low-level layer primitives.

Convolution

CausalConv1d - Stateful causal 1D convolution for recurrent use.

ParallelCausalConv1d - Parallel (non-recurrent) causal convolution.

Dense

BlockDiagonalDense - Block-diagonal dense layer for efficient computation.

Normalization

MultiHeadLayerNorm - Per-head layer normalization.

Utility

Flatten - Reshape to batch × feature.

Identity - Pass-through layer.