Nn_blocks (neural_nets_lib.Ocannl.Nn

Neural Network Building Blocks

This file contains basic building blocks for neural networks, with limited functionality. Feel free to copy-paste and modify as needed.

Design principles, OCANNL fundamentals, and common patterns:

"Principle of least commitment": use row variables where axis count doesn't matter
Einsum specs here often use single-char mode (no commas) but with spaces for readability
Pooling uses constant kernels (0.5 + 0.5) to propagate window dimensions
conv2d uses convolution syntax: "stride*out+kernel," (often in multi-char mode)
Input axes (before →) for kernels show intent (and end up rightmost for memory locality)
Inline params are always learnable and are lifted to unit parameter ()
Introduce inputs to a block after sub-block construction (sub-blocks have no automatic lifting like there is for inline definitions of params)
Always use literal strings with einsum operators when capturing variables
Avoid unnecessary variable captures in einsum operators, be mindful they can shadow other identifiers

module Tn = Ocannl.Operation.DSL_modules.Ir.Tnode

val mlp_layer : 
  label:Base.string Base.list ->
  hid_dim:Base.int ->
  unit ->
  Tensor.t ->
  Tensor.t

val dropout : 
  rate:Base.Float.t ->
  unit ->
  train_step:Ir.Indexing.static_symbol option ->
  Tensor.t ->
  Tensor.t

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.

val mlp : 
  label:Base.string Base.list ->
  hid_dims:Base.int Base.List.t ->
  unit ->
  Tensor.t ->
  Tensor.t

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

val reduce_specified_axes : Base.String.t -> Base.String.t

val softmax : 
  spec:Base.String.t ->
  ?temperature:Base.Float.t ->
  unit ->
  Tensor.t ->
  Tensor.t

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

val multi_head_attention : 
  label:Base.string Base.list ->
  num_heads:Base.int ->
  ?temperature:Base.Float.t ->
  ?dropout_rate:Base.Float.t ->
  unit ->
  train_step:Ir.Indexing.static_symbol option ->
  ?mask:Tensor.t ->
  Tensor.t ->
  Tensor.t

val layer_norm : 
  label:Base.string Base.list ->
  ?epsilon:Base.float ->
  unit ->
  Tensor.t ->
  Tensor.t

val transformer_encoder_block : 
  label:Base.string Base__List.t ->
  num_heads:Base.int ->
  d_ff:Base.int ->
  ?epsilon:Base.float ->
  unit ->
  train_step:Ir.Indexing.static_symbol option ->
  Tensor.t ->
  Tensor.t

val cross_attention : 
  label:Base.string Base.list ->
  num_heads:Base.int ->
  ?temperature:Base.Float.t ->
  ?dropout_rate:Base.Float.t ->
  unit ->
  train_step:Ir.Indexing.static_symbol option ->
  Tensor.t ->
  enc_output:Tensor.t ->
  Tensor.t

val transformer_decoder_block : 
  label:Base.string Base__List.t ->
  num_heads:Base.int ->
  d_ff:Base.int ->
  ?epsilon:Base.float ->
  unit ->
  train_step:Ir.Indexing.static_symbol option ->
  Tensor.t ->
  enc_output:Tensor.t ->
  mask:Tensor.t ->
  Tensor.t

val transformer_encoder : 
  label:Base.String.t Base__List.t ->
  num_layers:int ->
  num_heads:Base.int ->
  d_ff:Base.int ->
  ?epsilon:Base.float ->
  unit ->
  train_step:Ir.Indexing.static_symbol option ->
  Tensor.t ->
  Tensor.t

val transformer_decoder : 
  label:Base.String.t Base__List.t ->
  num_layers:int ->
  num_heads:Base.int ->
  d_ff:Base.int ->
  ?epsilon:Base.float ->
  unit ->
  train_step:Ir.Indexing.static_symbol option ->
  Tensor.t ->
  enc_output:Tensor.t ->
  mask:Tensor.t ->
  Tensor.t

val transformer : 
  label:Base.string Base.list ->
  num_encoder_layers:int ->
  num_decoder_layers:int ->
  num_heads:Base.int ->
  d_model:Base.int ->
  d_ff:Base.int ->
  ?epsilon:Base.float ->
  unit ->
  train_step:Ir.Indexing.static_symbol option ->
  src:Tensor.t ->
  tgt:Tensor.t ->
  mask:Tensor.t ->
  Tensor.t

Convolutional Neural Network Building Blocks

val conv2d : 
  label:Base.string Base.list ->
  ?kernel_size:Base.int ->
  ?stride:Base.Int.t ->
  ?use_padding:Base.bool ->
  unit ->
  Tensor.t ->
  Tensor.t

2D convolution layer with flexible padding and stride options.

val depthwise_separable_conv2d : 
  label:Base.string Base.list ->
  ?kernel_size:Base.int ->
  ?stride:Base.Int.t ->
  unit ->
  Tensor.t ->
  Tensor.t

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)

val max_pool2d : 
  ?stride:Base.Int.t ->
  ?window_size:Base.int ->
  unit ->
  Tensor.t ->
  Tensor.t

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

val avg_pool2d : 
  ?stride:Base.Int.t ->
  ?window_size:Base.int ->
  unit ->
  Tensor.t ->
  Tensor.t

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

val global_avg_pool2d : Ocannl.Operation.DSL_modules.Tensor.t -> Tensor.t

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

val batch_norm2d : 
  label:Base.string Base.list ->
  ?epsilon:Base.float ->
  ?momentum:float ->
  unit ->
  train_step:'a option ->
  Tensor.t ->
  Tensor.t

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

val conv_bn_relu : 
  label:Base.string Base__List.t ->
  ?kernel_size:Base.int ->
  ?stride:Base.Int.t ->
  unit ->
  train_step:'a option ->
  Tensor.t ->
  Tensor.t

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

val resnet_block : 
  label:Base.string Base__List.t ->
  ?stride:Base.Int.t ->
  unit ->
  train_step:'a option ->
  Tensor.t ->
  Tensor.t

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

val lenet : 
  label:Base.string Base.list ->
  ?num_classes:Base.int ->
  unit ->
  train_step:'a ->
  Tensor.t ->
  Tensor.t

LeNet-style architecture for simple image classification (e.g., MNIST). Classic architecture: conv -> pool -> conv -> pool -> fc layers

val vgg_block : 
  label:Base.string Base__List.t ->
  num_convs:int ->
  ?kernel_size:Base.int ->
  unit ->
  train_step:'a option ->
  Tensor.t ->
  Tensor.t

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

val sokoban_cnn : 
  label:Base.string Base.list ->
  ?num_actions:Base.int ->
  unit ->
  train_step:'a option ->
  grid_state:Tensor.t ->
  Tensor.t * Tensor.t

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

val mobile_cnn : 
  label:Base.string Base.list ->
  ?num_classes:Base.int ->
  ?width_mult:float ->
  unit ->
  train_step:'a option ->
  Tensor.t ->
  Tensor.t

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