conv1d ~in_channels ~out_channels ?kernel_size ?stride ?dilation ?padding () creates a 1D convolutional layer over inputs of shape batch; in_channels; length. Supports `Same, `Valid, and `Causal padding. Default: kernel_size=3, stride=1, dilation=1, padding=`Same.

conv2d ~in_channels ~out_channels ?kernel_size () creates a 2D convolutional layer.

Performs 2D convolution over 4D input tensors of shape batch_size, in_channels, height, width. The layer maintains learnable weight and bias parameters.

• parameter in_channels

  Number of input channels

• parameter out_channels

  Number of output filters

• parameter kernel_size

  Filter dimensions as (height, width). Default: (3, 3)

The weight tensor has shape out_channels, in_channels, kernel_height, kernel_width and is initialized using Glorot uniform initialization. The bias tensor has shape out_channels and is zero-initialized.

Example

  let conv = Layer.conv2d ~in_channels:3 ~out_channels:64 ~kernel_size:(5, 5) () in
  (* Processes RGB images (3 channels) to produce 64 feature maps with 5x5 filters *)

linear ~in_features ~out_features ?weight_init ?bias_init () creates a fully connected layer.

Applies linear transformation y = xW^T + b where x is input, W is weight matrix, and b is bias vector. Accepts inputs of any shape with last dimension matching in_features.

• parameter in_features

  Size of input feature dimension

• parameter out_features

  Size of output feature dimension

• parameter weight_init

  Weight initialization strategy. Default: Initializers.glorot_uniform

• parameter bias_init

  Bias initialization strategy. Default: Initializers.zeros

The weight tensor has shape out_features, in_features and bias has shape out_features.

Examples

  let classifier = Layer.linear ~in_features:512 ~out_features:10 () in
  (* Maps 512-dimensional features to 10 class logits *)

  let custom_init = Layer.linear
    ~in_features:256 ~out_features:128
    ~weight_init:(Initializers.he_normal ())
    ~bias_init:(Initializers.constant 0.1) () in

dropout ~rate () creates a dropout layer for regularization.

During training, randomly sets elements to zero with probability rate and scales remaining elements by 1 / (1 - rate) to maintain expected values. During evaluation, applies identity transformation.

• parameter rate

  Dropout probability in range 0.0, 1.0

Requires random number generator during training. No learnable parameters.

Example

  let drop = Layer.dropout ~rate:0.5 () in
  (* Randomly zeros 50% of activations during training *)

batch_norm ~num_features () creates a batch normalization layer.

Normalizes inputs across the batch dimension, learning scale and shift parameters. Applies transformation y = γ((x - μ) / σ) + β where μ and σ are batch statistics, and γ, β are learnable parameters.

• parameter num_features

  Number of features to normalize (typically channel dimension)

Maintains running statistics for evaluation mode. Parameters include scale (γ), bias (β), running mean, and running variance.

max_pool2d ~kernel_size ?stride () creates a 2D max pooling layer.

Applies maximum operation over spatial windows, reducing spatial dimensions while preserving channel dimension.

• parameter kernel_size

  Pooling window size as (height, width)

• parameter stride

  Pooling stride as (height, width). Default: same as kernel_size

No learnable parameters.

avg_pool2d ~kernel_size ?stride () creates a 2D average pooling layer.

Applies average operation over spatial windows, providing smoother downsampling compared to max pooling.

• parameter kernel_size

  Pooling window size as (height, width)

• parameter stride

  Pooling stride as (height, width). Default: same as kernel_size

No learnable parameters.

flatten () creates a flatten layer that reshapes multidimensional inputs to 2D.

Preserves batch dimension while flattening all other dimensions. Transforms shape batch_size, d1, d2, ..., dn to batch_size, d1 * d2 * ... * dn.

Commonly used before dense layers in CNN architectures. No learnable parameters.

relu () creates a ReLU activation layer applying max(0, x) elementwise.

Most common activation for hidden layers. Computationally efficient with good gradient flow for positive inputs. No learnable parameters.

sigmoid () creates a sigmoid activation layer applying 1 / (1 + exp(-x)) elementwise.

Maps inputs to range (0, 1). Commonly used for binary classification and gating mechanisms. No learnable parameters.

tanh () creates a hyperbolic tangent activation layer applying tanh(x) elementwise.

Maps inputs to range (-1, 1). Provides stronger gradients than sigmoid but can suffer from vanishing gradients. No learnable parameters.

gelu () creates a GELU activation layer.

Applies Gaussian Error Linear Unit activation, popular in transformer architectures. Smoother alternative to ReLU with better gradient properties. No learnable parameters.

swish () creates a Swish activation layer applying x * sigmoid(x) elementwise.

Self-gated activation function that can outperform ReLU in deep networks. No learnable parameters.

sequential layers creates a sequential composition of layers.

Applies layers in order, threading output of each layer as input to the next. The resulting module's parameters are the union of all component layer parameters.

• parameter layers

  List of layers to compose

Example

  let mlp = Layer.sequential [
    Layer.linear ~in_features:784 ~out_features:256 ();
    Layer.relu ();
    Layer.dropout ~rate:0.3 ();
    Layer.linear ~in_features:256 ~out_features:10 ();
  ] in

einsum ~einsum_str ~shape ?kernel_init () creates a parameterized Einstein summation layer.

Implements learnable tensor contractions specified by Einstein notation. Useful for implementing custom linear transformations and attention mechanisms.

• parameter einsum_str

  Einstein summation string describing the contraction

• parameter shape

  Shape of the learnable kernel parameter

• parameter kernel_init

  Kernel initialization strategy. Default: Initializers.glorot_uniform

rms_norm ~dim ?eps ?scale_init () creates a Root Mean Square normalization layer.

Applies RMS normalization with learnable scaling. Normalizes by the RMS of activations rather than full statistics like batch normalization.

• parameter dim

  Dimension to normalize over

• parameter eps

  Small constant for numerical stability. Default: 1e-6

• parameter scale_init

  Scale parameter initialization. Default: Initializers.ones

layer_norm ~dim ?eps ?elementwise_affine () creates a layer normalization layer.

Normalizes activations across the feature dimension within each sample. Popular in transformer architectures for stable training.

• parameter dim

  Dimension to normalize over

• parameter eps

  Small constant for numerical stability. Default: 1e-6

• parameter elementwise_affine

  Whether to learn scale and shift parameters. Default: true

embedding ~vocab_size ~embed_dim ?scale ?embedding_init () creates an embedding lookup layer.

Maps discrete tokens (integers) to dense vectors. Commonly used as the first layer in NLP models to convert token IDs to continuous representations.

• parameter vocab_size

  Size of the vocabulary (number of possible tokens)

• parameter embed_dim

  Dimensionality of embedding vectors

• parameter scale

  Whether to scale embeddings by sqrt(embed_dim). Default: false

• parameter embedding_init

  Embedding matrix initialization. Default: Initializers.normal

The embedding matrix has shape vocab_size, embed_dim.

mlp ~in_features ~hidden_features ~out_features ... creates a multi-layer perceptron (feed-forward network).

Standard MLP architecture: Linear -> Activation -> Dropout -> Linear -> Dropout Commonly used in transformers and other architectures.

• parameter in_features

  Input dimension

• parameter hidden_features

  Hidden layer dimension

• parameter out_features

  Output dimension

• parameter activation

  Activation function. Default: `gelu

• parameter dropout

  Dropout probability. Default: 0.0

Simple tanh RNN over a sequence. Input batch; seq; input_size, output batch; hidden_size (last hidden state).

GRU over a sequence. Input/output like rnn.

LSTM over a sequence. Input/output like rnn.

Adds learned positional embeddings to input batch; seq; embed_dim.

Create a max_len; embed_dim sinusoidal positional encoding table (not trainable). Can be added to token embeddings.