Module Ocannl.Nn_blocks

Convert a list of integers to a compact tensor of class IDs (no num_classes allocation).

• parameter lst

  List of integer class indices (0-based)

• returns

  A tensor of shape len (a len-sized batch axis) holding the IDs in uint32 precision, so that the gh-343 embedding gather can guard the dynamic index with native integer comparisons (no double-precision guard, no integrality check).

Build a logical one-hot tensor from a tensor of class IDs, using only existing operations (range + equality) so the compiler keeps the proof that the result is one-hot (enabling the gh-343 embedding gather optimization). No dense len * num_classes data is materialized on the host. With ids shaped as a len batch (output rank 0), the result is len; num_classes: one_hot[i, k] = (k == ids[i]).

• parameter num_classes

  The number of classes (size of the one-hot dimension).

Convert a list of integers to a logical one-hot encoded tensor of shape len; num_classes. This composes class_ids_of_int_list and one_hot_of_ids: it stores only len compact IDs on the host and expresses the one-hot logically, rather than allocating a dense len * num_classes Bigarray. See dense_one_hot_of_int_list if a materialized host one-hot is genuinely required.

• parameter num_classes

  The number of classes (size of the one-hot dimension)

• parameter lst

  List of integer class indices (0-based)

Convert a list of integers to a dense, host-materialized one-hot Bigarray-backed tensor of shape len; num_classes. Prefer one_hot_of_int_list (logical) unless a dense host fixture is needed; a materialized Bigarray carries no proof that it is one-hot, so it cannot be optimized into an embedding gather.

Convert an array of token IDs (e.g. the output of Dataprep.Bpe.encode) to a tensor of shape len (a len-sized batch axis, output rank 0) holding the IDs in uint32 precision -- large enough for any practical vocabulary (e.g. Gemma's 256K) and integer-native for the gh-343 embedding gather (see class_ids_of_int_list). Feed the result to one_hot_of_ids for an embedding lookup.

If max_len is given, the sequence is truncated or right-padded with pad_id to exactly max_len (needed for fixed-shape batched inference). Prepending/appending special tokens (BOS/EOS) is the caller's responsibility.

• parameter pad_id

  The token ID used for padding (default 0)

• parameter max_len

  If given, the result has exactly this length regardless of Array.length ids

Batched variant of token_ids_of_array: convert several token-ID sequences to a single tensor of shape num_seqs; max_len (two batch axes, output rank 0) in uint32 precision. Each sequence is truncated or right-padded with pad_id to max_len, which defaults to the length of the longest sequence. Composes with one_hot_of_ids the same way as the unbatched variant, giving a num_seqs; max_len; num_classes logical one-hot.

• parameter pad_id

  The token ID used for padding (default 0)

• parameter max_len

  The common sequence length (default: longest sequence in seqs)

Gaussian Error Linear Unit, tanh approximation (Hendrycks & Gimpel 2016): 0.5 * x * (1 + tanh (sqrt (2/pi) * (x + 0.044715 * x^3))). This is the GPT-2 activation (HF gelu_new); it uses only existing primitives (notably Tanh_approx), no dedicated backend support needed.

Masks and scales by 1/keep_prob to maintain expected value. When train_step = None, the dropout rate is ignored and the tensor is returned unmodified.

Multi-layer perceptron of depth List.length hid_dims + 1, with a linear output layer.

Softmax across specified axes. Does not support non-default row variables.

Cross-entropy loss from raw logits, adopting the llm.c fused_classifier contract (issue gh-464): numerically stable log-sum-exp (never log (softmax logits)), written so that the softmax-probabilities intermediate stays virtual — never materialized in either pass — and the backward accumulates probs - targets directly into the logits gradient. The row-wise max is detached via stop_gradient: the max term cancels exactly in the loss, so no gradient should flow through it (llm.c likewise treats it as a constant).

spec follows the softmax convention: it names the class/vocabulary axes to reduce over (e.g. "... | v"). targets is a probability distribution over those axes (typically one-hot; label smoothing works too). mask multiplies the per-position losses (e.g. 0/1 to exclude padding tokens); normalize_by divides the summed loss (e.g. by a valid-token count). Without mask and normalize_by, returns the summed scalar loss.

Strategy for positional encoding in attention / transformer blocks.

RoPE inverse frequencies: theta_k = base^(-2k/d) for k = 0..half_d-1.

• parameter half_d

  Per-head key dimension divided by 2 (i.e. d_k / 2, NOT d_model / 2).

• parameter base

  Default 10000.0; some models use 500000 for long contexts.

Position indices 0, 1, ..., seq_len-1 as a non-learned batch-dim tensor.

Sinusoidal positional encoding (Vaswani et al. 2017). Non-learned, shape: batch_dims=max_len, output_dims=d_model. Matches model width at the transformer input level, NOT per-head width.

Apply RoPE rotation to tensor x whose last output axis has even size d. Rotates within the last output axis (per-head width d) without crossing head boundaries. freqs has output=d/2, positions has batch=seq_len.

Decoder-only transformer block: masked self-attention + FFN with post-norm LayerNorm. Like transformer_encoder_block but accepts a ~mask parameter for causal masking. No cross-attention — suitable for autoregressive language models.

Stack of decoder_only_block layers.

Transformer with teacher forcing for autoregressive training.

TODO: Simplify once tensor shifting/slicing is better supported in shape inference. Currently requires pre-shifted tgt_input (all but last token) and tgt_target (all but first token). During training, the model learns to predict tgt_target given tgt_input.

2D convolution layer with flexible padding and stride options.

When use_padding=false and stride > 1, the input spatial dimensions must satisfy: (input_size - kernel_size) mod stride = 0, otherwise shape inference will fail with "incompatible stride" error. The output size is (input_size - kernel_size) / stride + 1.

When use_padding=true, there is no such restriction and output size is input_size / stride.

• parameter out_channels

  Optional number of output channels. If not provided, must be inferred from context (e.g., from a downstream operation that constrains the output shape).

Depthwise separable convolution - more efficient for mobile/edge devices. Consists of depthwise conv (spatial filtering per channel) followed by pointwise conv (1x1 conv for channel mixing).

See conv2d for dimension constraints when use_padding=false.

Max pooling for 2D spatial data - reduces spatial dimensions by taking maximum values.

The input spatial dimensions must satisfy: (input_size - window_size) mod stride = 0, otherwise shape inference will fail. The output size is (input_size - window_size) / stride + 1.

Note: The < in the einsum spec indicates no-padding mode (indices stay within bounds).

Average pooling for 2D spatial data - reduces spatial dimensions by averaging values.

See max_pool2d for dimension constraints.

Global average pooling - reduces each feature map to a single value by averaging. Commonly used before final classification layer.

Batch normalization for CNN layers - normalizes across the batch dimension for each channel. Typically applied after convolutions and before activations.

Batch normalization for MLP layers - normalizes across the batch axis only. Unlike batch_norm2d there are no spatial axes to reduce over; channel axes are carried through unchanged via the ..c.. row variable.

See the FIXME on batch_norm2d: running statistics are not implemented, so momentum is ignored and inference falls back to the learned gamma/beta parameters rather than population statistics. Acceptable for tutorial examples; do not rely on inference correctness for distribution-shifted inputs.

Conv block with conv -> batch norm -> activation pattern

Residual block for ResNet-style architectures. Features skip connections that help with gradient flow in deep networks.

LeNet-style architecture for simple image classification (e.g., MNIST). Classic architecture: conv -> pool -> conv -> pool -> fc layers. Output shape is inferred from training data.

VGG-style block - multiple convolutions with same filter count followed by pooling

Simple CNN for Sokoban-like grid environments. Processes grid states with multiple conv layers and outputs action logits.

Modern CNN with depthwise separable convolutions for efficiency. Suitable for mobile/edge deployment.