Building Neural Networks in Clorch: Layers and Lifecycle

Complete reference for clorch.nn: linear, conv, recurrent, normalization, pooling, defmodel macro, lifecycle API, and model inspection with nn/summary.

Clorch’s clorch.nn namespace is the primary entry point for building neural networks. Every layer wraps LibTorch’s C++ frontend directly, giving you native performance with no intermediate overhead — while remaining composable with ordinary Clojure data structures like vectors and maps.

Common Layers

Linear (Dense) Layers

nn/linear creates a fully-connected layer backed by LibTorch’s LinearImpl. The optional :bias false keyword disables the bias term.

(require '[clorch.nn :as nn])

;; Linear(in_features=10, out_features=5)
(nn/linear 10 5)

;; Linear without bias
(nn/linear 10 5 :bias false)

Convolutional Layers

Clorch exposes 1D, 2D, and 3D convolutions and their transpose variants with the full set of spatial options.

;; Conv2d(in_channels=3, out_channels=16, kernel_size=3)
(nn/conv2d 3 16 3)

;; Advanced parameters: rectangular kernel, stride, padding, dilation, groups
(nn/conv2d 3 16 [3 5] :stride 2 :padding 1 :dilation 2 :groups 1)

;; 1D and 3D convolutions share the same option set
(nn/conv1d 1 16 3 :stride 1 :padding 1)
(nn/conv3d 1 8 3 :stride 1 :padding 1)

;; Transpose Convolution (decoder/upsampling)
(nn/conv-transpose2d 16 3 3 :stride 2)
(nn/conv-transpose1d 16 3 3 :stride 2)
(nn/conv-transpose3d 16 3 3 :stride 2)

Recurrent Layers

LSTM, GRU, and plain RNN delegate directly to LibTorch’s optimized C++ RNN implementations.

;; LSTM(input_size=10, hidden_size=20)
(nn/lstm 10 20)

;; Bidirectional LSTM with multiple layers and dropout
(nn/lstm 10 20 :num-layers 2 :bidirectional true :dropout 0.1)

;; Bidirectional GRU with multiple layers and dropout
(nn/gru 10 20 :num-layers 2 :bidirectional true :dropout 0.1)

;; Plain RNN
(nn/rnn 10 20 :num-layers 1 :batch-first true)

Normalization

(nn/batchnorm1d 64)          ;; Batch Normalization for 1D (sequence) feature maps
(nn/batchnorm2d 64)          ;; Batch Normalization for 2D (image) feature maps
(nn/batchnorm3d 64)          ;; Batch Normalization for 3D (volumetric) feature maps
(nn/layernorm 128)            ;; Layer Normalization (custom Clorch record)
(nn/rmsnorm 512)              ;; High-performance RMSNorm for LLMs
(nn/groupnorm 8 64)           ;; Group Normalization: 8 groups, 64 channels
(nn/instancenorm1d 64)        ;; Instance Normalization (1D)
(nn/instancenorm2d 64)        ;; Instance Normalization (2D)
(nn/instancenorm3d 64)        ;; Instance Normalization (3D)

nn/layernorm and nn/rmsnorm are implemented as Clojure records (not native modules) so their learnable parameters participate fully in nn/parameters, nn/to, and state-dict traversal.

Pooling & Padding

;; Max and average pooling
(nn/max-pool2d 2)
(nn/avg-pool2d 2)
(nn/adaptive-avg-pool2d [7 7])

;; Zero padding
(nn/zeropad2d 1)

;; Reflection and constant padding
(nn/reflection-pad2d 2)
(nn/constant-pad2d 1 3.14)

Containers

nn/sequential returns a plain Clojure vector. Because Clorch’s IModule protocol is extended to APersistentVector, the vector itself acts as a sequential container: nn/forward, nn/train, and nn/to all traverse it automatically.

(def model (nn/sequential
             (nn/linear 10 20)
             (nn/relu)
             (nn/dropout 0.5)
             (nn/linear 20 1)))

Embeddings

;; Learnable embedding table
(nn/embedding 10000 256)

;; Initialize from pre-trained weight matrix (frozen by default)
(nn/embedding-from-pretrained pretrained-weight)

;; Allow the embeddings to be fine-tuned
(nn/embedding-from-pretrained pretrained-weight :freeze false)

Utility Layers

;; Flatten all dimensions starting from dim 1
(nn/flatten)

;; Unflatten dimension 1 into [2 8]
(nn/unflatten 1 [2 8])

;; Pass-through identity layer
(nn/identity)

;; Nearest-neighbour upsample to a fixed size or by a scale factor
(nn/upsample :size [64 64] :mode :nearest)
(nn/upsample :scale-factor 2 :mode :nearest)

---

Custom Models with defmodel

defmodel is a macro that generates a fully-featured Clojure record implementing the IModule protocol. It eliminates boilerplate: no manual defrecord, no hand-written -train or -to implementations.

The Three Parts

A defmodel form has exactly three parts:

1. Constructor arguments — the parameters your model factory function accepts.
2. Binding vector — field/value pairs evaluated once when the model is instantiated, identical in structure to let.
3. forward form — a method body that may reference any binding by name.

(require '[clorch.torch :as t]
         '[clorch.nn :as nn]
         '[clorch.nn.functional :as F])

(nn/defmodel MyClassifier [in-dim hidden-dim num-classes]
  [l1 (nn/linear in-dim hidden-dim)
   l2 (nn/linear hidden-dim num-classes)]
  (forward [x]
    (nn/forward l2 (F/relu (nn/forward l1 x)))))

Constructing and Calling the Model

;; Instantiate: calls the generated constructor function
(def classifier (MyClassifier 784 256 10))

;; Explicit forward pass
(def logits (nn/forward classifier (t/randn [32 784])))
(t/size logits) ; => [32 10]

;; Callable syntax — the generated record implements clojure.lang.IFn
(classifier (t/randn [32 784]))

How Registered Fields Participate in Lifecycle Operations

Any field whose value is a native Module, a Parameter, a Tensor, a TensorVector, a Clojure vector, or a Clojure map is automatically traversed by every lifecycle operation:

Operation	Behaviour on fields
nn/train	Recursively sets training/eval mode on all sub-modules
nn/to	Recursively moves tensors and parameters to device/dtype
nn/parameters	Recursively collects all leaf tensors into a TensorVector
nn/state-dict	Recursively extracts a nested map of weight tensors
nn/load-state-dict	Recursively copies tensors back into the model

---

Model Inspection with nn/summary

nn/summary prints a PyTorch-style table showing the output shape and parameter count of every layer visited during a forward pass.

Usage

Simple shape input

Pass a shape vector and Clorch will synthesize a randn tensor of that shape automatically:

(nn/summary model [1 784])

Complex / LLM input

Pass an explicit tensor or map when your model requires multiple arguments or integer tensors:

(nn/summary gpt {:idx  (t/zeros [8 4] {:dtype :int64})
                 :mask (t/ones [4 4])})

How the Dry-Run Works

When you call nn/summary, Clorch:

1. Sets the model to eval mode temporarily.
2. Wraps execution in autograd/no-grad to suppress gradient bookkeeping.
3. Binds the dynamic var *trace* to an atom that intercepts every nn/forward call.
4. Captures the module type, output shape, and parameter count of every layer encountered during that single pass.

Because the trace is driven by a real forward pass, it faithfully captures dynamic shapes, conditional branches, and any reshape operations inside your model logic.

Sample Output

----------------------------------------------------------------
Layer (type)                   Output Shape         Param #
================================================================
Linear                         [1 20]               220
ReLU                           [1 20]               0
Linear                         [1 5]                105
PersistentVector               [1 5]                325
================================================================
Total params: 325
Trainable params: 325
Non-trainable params: 0
----------------------------------------------------------------

The outermost row (PersistentVector) represents the top-level sequential container. Its parameter count is the cumulative total of all layers inside it.

---

Lifecycle API

Mode Management

(nn/train model true)   ;; Training mode — enables dropout and batchnorm updates
(nn/train model false)  ;; Evaluation mode — freezes batchnorm, disables dropout

Device and Dtype Transfer

nn/to recursively traverses the entire model tree — native modules, records, vectors, maps, and bare tensors — and moves every parameter to the target device or dtype.

(nn/to model :cuda)     ;; Move entire model to GPU
(nn/to model :float16)  ;; Cast all parameters to half precision
(nn/to model :float64)  ;; Cast all parameters to double

Parameters and State Dictionaries

;; Collect all leaf parameters into a native TensorVector (for optimizers)
(nn/parameters model)

;; Extract a nested Clojure map of all named weights and biases
(nn/state-dict model)

;; Copy weights from a saved state dict into an existing model
(nn/load-state-dict model saved-sd)

;; Save model weights to disk (state-dict format)
(nn/save-weights model "checkpoint.pt")

;; Load weights from disk into an existing model
(nn/load-weights model "checkpoint.pt")

Gradient Management

;; Zero all gradients in the model (alternative to optim/zero-grad on the model)
(nn/zero-grad model)

;; Clip gradient norm to prevent exploding gradients
(nn/clip-grad-norm! (nn/parameters model) 1.0)

Forward Pass

(nn/forward model input)

nn/forward is polymorphic: it accepts native Module instances, defmodel records, plain Clojure vectors (sequential), and Clojure functions.

Introspection

;; Flat map of all parameter name paths to tensors
(nn/named-parameters model)

;; Sequence of all sub-modules (native and record) in the tree
(nn/modules model)

;; Apply a side-effecting function to every module in the tree
(nn/apply model (fn [m] (println (type m))))

---

LLM-Specific Modules

Clorch ships with production-ready building blocks for large language models, all defined using defmodel.

RMSNorm

Root Mean Square Layer Normalization — lighter than LayerNorm as it omits the mean-centering step.

(nn/rmsnorm 512)              ;; 512-dimensional hidden size
(nn/rmsnorm 512 :eps 1e-6)    ;; Custom epsilon

SwiGLU

Gated feed-forward block used in LLaMA-style models. The gate is computed with SiLU.

(nn/defmodel SwiGLU [dim hidden-dim]
  [w1 (nn/linear dim hidden-dim)
   w2 (nn/linear dim hidden-dim)
   w3 (nn/linear hidden-dim dim)]
  (forward [x]
    (let [gate (F/silu (nn/forward w1 x))
          feat (nn/forward w2 x)]
      (nn/forward w3 (t/mul gate feat)))))

GroupedQueryAttention

Multi-head attention with Grouped Query Attention (GQA), RoPE embeddings, and KV-cache support. Delegates to torch/scaled_dot_product_attention for fused CUDA dispatch.

;; 512-dim, 8 query heads, 2 KV heads, 2048 context length, 0.1 dropout
(GroupedQueryAttention 512 8 2 2048 0.1)

Pass input as a map to supply optional components:

(nn/forward attn {:x       hidden-states
                  :mask    causal-mask
                  :freqs   [cos-emb sin-emb]
                  :kv-cache kv-cache-atom})

nn/generate — Autoregressive Token Generation

Generates tokens one step at a time using multinomial sampling, with automatic context-window truncation.

;; model: any defmodel with a forward that accepts token indices
;; idx: [batch, seq-len] int64 tensor of prompt tokens
(nn/generate model idx 100 context-size)