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.

The row max is taken through the non-differentiable DSL: softmax is invariant under a shift of its input, so the gradient through the max is exactly zero, and the composed autodiff would otherwise spend a per-row reduction and an argmax-equality pass computing that zero up to rounding -- and inject the rounding at the argmax cells. Applying an NTDSL operation to a differentiable operand cuts the gradient at that operation (Prohibit_grad); stop_gradient is the identity instance of the same mechanism, at the price of a copy.

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 taken through the non-differentiable DSL, as in softmax: 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.

The score a masked-out attention position is filled with before the softmax: Float.neg_infinity, so that exp (fill - max) is exactly zero.

This is deliberately not a large finite magic number. A finite sentinel has to be chosen against a storage precision — the long-standing -1e9 is four orders of magnitude past fp16's largest finite value, so it tripped Ops.exceeds_fp16_cutoff and made every masked attention unlowerable at half precision. Negative infinity is exactly representable in every float format OCANNL supports, and says what is meant: this position is below anything the format can hold.

The one behavior a finite fill bought is that a fully masked row (no live position at all) degrades to a uniform distribution rather than to NaN — with an all--inf row the softmax's max-subtraction yields -inf - -inf = nan. Causal masks always leave the diagonal live and are unaffected; a padding mask that can mask a whole row should pass a finite ?mask_fill (e.g. -1e4, which still underflows exp to zero at every precision).

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.

Without padding (use_padding=false, the default), 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. The < in the einsum spec indicates no-padding mode (indices stay within bounds).

With use_padding=true ("same" pooling, output size input_size / stride), the pool lowers with clamped window bounds (gh-504): per output position the window loop is range-guarded to the intersection of the window with the valid input range, and an out-of-range position contributes the max-accumulation identity (-inf) — the same as not visiting it — so clamping is semantically exact and the operand needs NO margins at all. The operand therefore composes freely with 0-neutral margin-touching consumers such as padded convs (the Inception-block pattern), with eagerly allocated data nodes (e.g. TDSL.init), and with operands already lowered in earlier compilations of a staged flow. The backward argmax scatter transposes the clamp (guarded writes). Interior outputs still see full windows: the guards flip truth only at affine breakpoints of the output axis, so Schedule.Partition at Schedule.partition_breakpoints specializes guard-free interior segments (gh-508).

Overlapping pooling (stride < window_size, AlexNet-style) has exact gradients: the gradient gate lives in the (output x window) product space (gh-512), so each position receives gradient from exactly the windows it won, with ties gating every achieving pair. Non-overlapping pooling (stride >= window_size, the common case) dispatches to the cheaper input-space gate (gh-527) — exact on that domain, ties included, see Operation.tropical.

Like max_pool2d, but with use_padding=true the pool reads a private materialized copy of the operand instead of the operand itself. Since the clamped-window lowering (gh-504) removed the pool's -inf margin demand, max_pool2d composes with any operand and the copy is no longer needed as a conflict remedy — this variant is kept as the materialized-copy remedy pattern for the remaining conflict class (margin-touching consumers with different FINITE neutral elements, e.g. margins committed by wrap_padded to a value a padded conv's 0 neutral conflicts with), and for explicitly decoupling the pool's read from a shared buffer.

Prefer max_pool2d: it reads the operand directly, without the extra buffer and copy.

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

See max_pool2d for dimension constraints. Note there is no use_padding option: a padded (=-mode) add-family window keeps the physical 0-margins mechanism (finite accumulation identity — clamped lowering, gh-504, applies only to the max/tropical family), which gives count_include_pad semantics — the divisor is the full window size everywhere, counting the zero margins. PyTorch's default is count_exclude_pad (per-position valid-window divisor); adding a padded average pool is a per-op semantic decision that should document its choice.

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

Batch normalization for CNN layers, reducing batch and spatial axes per channel.

Some train_step normalizes with the batch mean and population variance and updates stored statistics after each execution. None uses the stored statistics without updating them. momentum defaults to 0.9 and is the retained fraction: running = momentum * running + (1 - momentum) * batch. Running mean/variance initialize to 0/1; gamma/beta initialize to 1/0. Statistics are non-trainable model state, initialized and persisted with the parameters. The variance uses the population divisor in both modes.

Batch normalization for MLP layers, reducing the batch axes and preserving channel axes. Training/inference behavior, initialization and the momentum recurrence are as in batch_norm2d.

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. Both convolutions use out_channels, or preserve the input channel row when it is omitted. An explicit out_channels selects a projection shortcut, including at stride 1; otherwise the shortcut is projected only when downsampling.

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. _width_mult is caller-visible as unimplemented: changing it currently has no effect on channel counts.