Ocannl.Nn_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:
are always learnable and are lifted to unit parameter ()module Tn = Ocannl_tensor.Operation.DSL_modules.Ir.Tnode
val mlp_layer :
label:Base.string Base.list ->
hid_dim:Base.int ->
unit ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val dropout :
rate:Base.Float.t ->
unit ->
train_step:Ir.Indexing.static_symbol option ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
Multi-layer perceptron of depth List.length hid_dims + 1
, with a linear output layer.
val softmax :
spec:Base.String.t ->
?temperature:Base.Float.t ->
unit ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
d_k:Base.int ->
d_v:Base.int ->
?temperature:Base.Float.t ->
?dropout_rate:Base.Float.t ->
unit ->
train_step:Ir.Indexing.static_symbol option ->
?mask:Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val layer_norm :
label:Base.string Base.list ->
?epsilon:Base.float ->
unit ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val transformer_encoder_block :
label:Base.string Base__List.t ->
num_heads:Base.int ->
d_k:Base.int ->
d_v:Base.int ->
d_ff:Base.int ->
?epsilon:Base.float ->
unit ->
train_step:Ir.Indexing.static_symbol option ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val cross_attention :
label:Base.string Base.list ->
num_heads:Base.int ->
d_k:Base.int ->
d_v:Base.int ->
?temperature:Base.Float.t ->
?dropout_rate:Base.Float.t ->
unit ->
train_step:Ir.Indexing.static_symbol option ->
Ocannl_tensor.Tensor.t ->
enc_output:Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val transformer_decoder_block :
label:Base.string Base__List.t ->
num_heads:Base.int ->
d_k:Base.int ->
d_v:Base.int ->
d_ff:Base.int ->
?epsilon:Base.float ->
unit ->
train_step:Ir.Indexing.static_symbol option ->
Ocannl_tensor.Tensor.t ->
enc_output:Ocannl_tensor.Tensor.t ->
mask:Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val transformer_encoder :
label:Base.String.t Base__List.t ->
num_layers:int ->
num_heads:Base.int ->
d_k:Base.int ->
d_v:Base.int ->
d_ff:Base.int ->
?epsilon:Base.float ->
unit ->
train_step:Ir.Indexing.static_symbol option ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val transformer_decoder :
label:Base.String.t Base__List.t ->
num_layers:int ->
num_heads:Base.int ->
d_k:Base.int ->
d_v:Base.int ->
d_ff:Base.int ->
?epsilon:Base.float ->
unit ->
train_step:Ir.Indexing.static_symbol option ->
Ocannl_tensor.Tensor.t ->
enc_output:Ocannl_tensor.Tensor.t ->
mask:Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val transformer :
label:Base.string Base.list ->
num_encoder_layers:int ->
num_decoder_layers:int ->
num_heads:Base__Int.t ->
d_enc:Base.int ->
d_dec:Base.int ->
d_ff:Base.int ->
?epsilon:Base.float ->
unit ->
train_step:Ir.Indexing.static_symbol option ->
src:Ocannl_tensor.Tensor.t ->
tgt:Ocannl_tensor.Tensor.t ->
mask:Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
val conv2d :
label:Base.string Base.list ->
?kernel_size:Base.int ->
?stride:Base.Int.t ->
?use_padding:Base.bool ->
unit ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
Average pooling for 2D spatial data - reduces spatial dimensions by averaging values.
val global_avg_pool2d :
Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.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:Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t * Ocannl_tensor.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 ->
Ocannl_tensor.Tensor.t ->
Ocannl_tensor.Tensor.t
Modern CNN with depthwise separable convolutions for efficiency. Suitable for mobile/edge deployment.