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 as floats.

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.

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.

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.