Skip to content

Commit 062ae64

Browse files
lukstaficlaude
lukstafi
and
claude
committed
Add comprehensive CNN building blocks and PyTorch/TF migration guide
- Add 2D convolutional layers with einsum notation (conv2d, depthwise_separable_conv2d) - Implement pooling operations (max_pool2d, avg_pool2d, global_avg_pool2d) - Add batch normalization for CNNs with train/inference modes - Create complete CNN architectures: - LeNet-style for MNIST-like tasks - ResNet blocks with skip connections - VGG-style blocks - Sokoban CNN for grid environments - MobileNet-style with depthwise separable convolutions - Add comprehensive migration guide from PyTorch/TensorFlow - Document OCANNL's unique approaches (no flattening needed, row variables) - Explain einsum notation modes (single-char vs multi-char) - Include common gotchas and idioms (0.5+0.5 trick, literal strings) Key design decisions: - Use row variables (..ic.., ..oc..) for flexible channel dimensions - Pooling uses constant kernels to carry shape info between inference phases - FC layers work directly with spatial dims (no flattening required) - Convolution syntax uses multi-char einsum mode with stride*out+kernel 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 49c1768 commit 062ae64

File tree

2 files changed

+314
-11
lines changed

2 files changed

+314
-11
lines changed

‎docs/migration_guide.md‎

Lines changed: 294 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,294 @@
1+
# Migration Guide: PyTorch/TensorFlow to OCANNL
2+
3+
This guide helps deep learning practitioners familiar with PyTorch or TensorFlow understand OCANNL's approach and idioms.
4+
5+
## Key Conceptual Differences
6+
7+
### Shape Inference vs Explicit Shapes
8+
- **PyTorch/TF**: Shapes are usually explicit (e.g., `Conv2d(in_channels=3, out_channels=64)`)
9+
- **OCANNL**: Shapes are inferred where possible, using row variables for flexibility
10+
```ocaml
11+
(* Channels as row variables allow multi-channel architectures *)
12+
conv2d ~label () x (* ..ic.. and ..oc.. are inferred *)
13+
```
14+
15+
### Two-Phase Inference System
16+
OCANNL separates shape inference from projection inference:
17+
- **Shape inference**: Global, propagates constraints across operations
18+
- **Projection inference**: Local per-assignment, derives loop structures from tensor shapes
19+
20+
This is why pooling needs a dummy constant kernel - to carry shape info between phases.
21+
22+
## Common Operations Mapping
23+
24+
| PyTorch/TensorFlow | OCANNL | Notes |
25+
|-------------------|---------|--------|
26+
| `x.view(-1, d)` or `x.reshape(-1, d)` | Not directly supported | Use manual dimension setting on constant tensor as workaround |
27+
| `x.flatten()` | Not supported | Future syntax might be: `"x,y => x&y"` |
28+
| `nn.Conv2d(in_c, out_c, kernel_size=k)` | `conv2d ~kernel_size:k () x` | Channels inferred or use row vars |
29+
| `F.max_pool2d(x, kernel_size=k)` | `max_pool2d ~window_size:k () x` | Uses `(0.5 + 0.5)` trick internally |
30+
| `F.avg_pool2d(x, kernel_size=k)` | `avg_pool2d ~window_size:k () x` | Normalized by window size |
31+
| `nn.BatchNorm2d(channels)` | `batch_norm2d () ~train_step x` | Channels inferred |
32+
| `F.dropout(x, p=0.5)` | `dropout ~rate:0.5 () ~train_step x` | Needs train_step for PRNG |
33+
| `F.relu(x)` | `relu x` | Direct function application |
34+
| `F.softmax(x, dim=-1)` | `softmax ~spec:"... \| ... -> ... d" () x` | Specify axes explicitly |
35+
| `torch.matmul(a, b)` | `a * b` or `a +* "...; ... => ..." b` | Einsum for complex cases |
36+
| `x.mean(dim=[1,2])` | `x ++ "... \| h, w, c => ... \| 0, 0, c" ["h"; "w"] /. (dim h *. dim w)` | Sum then divide |
37+
| `x.sum(dim=-1)` | `x ++ "... \| ... d => ... \| 0"` | Reduce by summing |
38+
39+
## Tensor Creation Patterns
40+
41+
### Parameters (Learnable Tensors)
42+
43+
| PyTorch | OCANNL |
44+
|---------|---------|
45+
| `nn.Parameter(torch.rand(d))` | `{ w }` or `{ w = uniform () }` |
46+
| `nn.Parameter(torch.randn(d))` | `{ w = normal () }` |
47+
| `nn.Parameter(torch.zeros(d))` | `{ w = 0. }` |
48+
| `nn.Parameter(torch.ones(d))` | `{ w = 1. }` |
49+
| With explicit dims | `{ w; o = [out_dim]; i = [in_dim] }` |
50+
51+
### Non-learnable Constants
52+
53+
| PyTorch | OCANNL | Notes |
54+
|---------|---------|--------|
55+
| `torch.ones_like(x)` | `0.5 + 0.5` | Shape-inferred constant 1 |
56+
| `torch.tensor(1.0)` | `!.value` or `1.0` | Scalar constant |
57+
| `torch.full_like(x, value)` | `NTDSL.term ~fetch_op:(Constant value) ()` | Shape-inferred |
58+
59+
## Network Architecture Patterns
60+
61+
### Sequential Models
62+
63+
**PyTorch:**
64+
```python
65+
model = nn.Sequential(
66+
nn.Conv2d(3, 64, 3),
67+
nn.ReLU(),
68+
nn.MaxPool2d(2),
69+
nn.Flatten(),
70+
nn.Linear(64*14*14, 10)
71+
)
72+
```
73+
74+
**OCANNL:**
75+
```ocaml
76+
let%op model () =
77+
let conv1 = conv2d ~kernel_size:3 () in
78+
let pool = max_pool2d () in
79+
fun x ->
80+
let x = conv1 x in
81+
let x = relu x in
82+
let x = pool x in
83+
(* No flattening needed - FC layer works with spatial dims *)
84+
{ w_out } * x + { b_out = 0.; o = [10] }
85+
```
86+
87+
### Residual Connections
88+
89+
**PyTorch:**
90+
```python
91+
class ResBlock(nn.Module):
92+
def __init__(self, channels):
93+
super().__init__()
94+
self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
95+
self.bn1 = nn.BatchNorm2d(channels)
96+
self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
97+
self.bn2 = nn.BatchNorm2d(channels)
98+
99+
def forward(self, x):
100+
identity = x
101+
out = self.conv1(x)
102+
out = self.bn1(out)
103+
out = F.relu(out)
104+
out = self.conv2(out)
105+
out = self.bn2(out)
106+
return F.relu(out + identity)
107+
```
108+
109+
**OCANNL:**
110+
```ocaml
111+
let%op resnet_block () =
112+
let conv1 = conv2d () in
113+
let bn1 = batch_norm2d () in
114+
let conv2 = conv2d () in
115+
let bn2 = batch_norm2d () in
116+
fun ~train_step x ->
117+
let identity = x in
118+
let out = conv1 x in
119+
let out = bn1 ~train_step out in
120+
let out = relu out in
121+
let out = conv2 out in
122+
let out = bn2 ~train_step out in
123+
relu (out + identity)
124+
```
125+
126+
## Einsum Notation
127+
128+
OCANNL's einsum is more general than PyTorch's, supporting row variables and convolutions.
129+
130+
### Syntax Modes
131+
132+
OCANNL's einsum has two syntax modes:
133+
134+
1. **Single-character mode** (PyTorch-compatible):
135+
- Triggered when NO commas appear in the spec
136+
- Each alphanumeric character is an axis identifier
137+
- Spaces are optional and ignored: `"ij"` = `"i j"`
138+
139+
2. **Multi-character mode**:
140+
- Triggered by ANY comma in the spec
141+
- Identifiers can be multi-character (e.g., `height`, `width`)
142+
- Must be separated by non-alphanumeric: `,` `|` `->` `;` `=>`
143+
- Enables convolution syntax: `stride*out+kernel`
144+
145+
| Operation | PyTorch einsum | OCANNL single-char | OCANNL multi-char |
146+
|-----------|---------------|-------------------|-------------------|
147+
| Matrix multiply | `torch.einsum('ij,jk->ik', a, b)` | `a +* "i j; j k => i k" b` | `a +* "i, j; j, k => i, k" b` |
148+
| Batch matmul | `torch.einsum('bij,bjk->bik', a, b)` | `a +* "b i j; b j k => b i k" b` | `a +* "batch, i -> j; batch, j -> k => batch, i -> k" b` |
149+
| Attention scores | `torch.einsum('bqhd,bkhd->bhqk', q, k)` | `q +* "b q \| h d; b k \| h d => b \| q k -> h" k` | `q +* "b, q \| h, d; b, k \| h, d => b \| q, k -> h" k` |
150+
| Convolution | N/A | N/A (needs multi-char) | `x +* "... \| stride*oh+kh, stride*ow+kw, ic; kh, kw, ic -> oc => ... \| oh, ow, oc" kernel` |
151+
152+
### Row Variables
153+
- `...` context-dependent ellipsis: expands to `..batch..` in batch position, `..input..` before `->`, `..output..` after `->`
154+
- `..b..` for batch axes (arbitrary number)
155+
- `..ic..`, `..oc..` for input/output channels (can be multi-dimensional)
156+
- `..spatial..` for spatial dimensions
157+
158+
## Common Gotchas and Solutions
159+
160+
### Variable Capture with Einsum
161+
❌ **Wrong:**
162+
```ocaml
163+
let spec = "... | h, w => ... | h0" in
164+
x ++ spec [ "h"; "w" ] (* Error: spec must be literal *)
165+
```
166+
167+
✅ **Right:**
168+
```ocaml
169+
x ++ "... | h, w => ... | h0" [ "h"; "w" ]
170+
```
171+
172+
### Creating Non-learnable Constants
173+
❌ **Wrong:**
174+
```ocaml
175+
{ kernel = 1. } (* Creates learnable parameter *)
176+
1.0 (* Creates fixed scalar shape *)
177+
```
178+
179+
✅ **Right:**
180+
```ocaml
181+
0.5 + 0.5 (* Both are shape-inferred constant 1 *)
182+
NTDSL.term ~fetch_op:(Constant 1.) ()
183+
```
184+
185+
### Parameter Scoping
186+
❌ **Wrong:**
187+
```ocaml
188+
let%op network () x =
189+
(* Sub-module defined after input *)
190+
let layer1 = my_layer () x in
191+
{ global_param } + x
192+
```
193+
194+
✅ **Right:**
195+
```ocaml
196+
let%op network () =
197+
(* Sub-modules before input *)
198+
let layer1 = my_layer () in
199+
fun x ->
200+
(* Inline definitions are lifted:
201+
used here, but defined before layer1 *)
202+
{ global_param } + layer1 x
203+
```
204+
205+
### Flattening for Linear Layers
206+
207+
⚠️ **Important:** OCANNL doesn't currently support flattening/reshaping operations.
208+
209+
```ocaml
210+
(* This performs REDUCTION (sum), not flattening: *)
211+
x ++ "... | ..spatial.. => ... | 0"
212+
213+
(* OCANNL's approach: Let FC layers work with multiple axes!
214+
Instead of flattening [batch, h, w, c] to [batch, h*w*c],
215+
just let your FC layer handle [batch, h, w, c] directly.
216+
The matrix multiplication will work across all the axes. *)
217+
218+
(* Example: FC layer after conv without flattening *)
219+
let%op fc_after_conv () x =
220+
(* x might have shape [batch, height, width, channels] *)
221+
{ w } * x + { b } (* w will adapt to match x's shape *)
222+
```
223+
224+
## Training Loop Patterns
225+
226+
### Basic Training Step
227+
228+
**PyTorch:**
229+
```python
230+
optimizer.zero_grad()
231+
output = model(input)
232+
loss = criterion(output, target)
233+
loss.backward()
234+
optimizer.step()
235+
```
236+
237+
**OCANNL (conceptual):**
238+
```ocaml
239+
(* OCANNL handles training differently - see Train module *)
240+
let sgd = Train.sgd_update ~learning_rate loss in
241+
let train_step = Train.to_routine ~ctx [%cd update; sgd] in
242+
(* Training happens via routines and contexts *)
243+
```
244+
245+
## Debugging Tips
246+
247+
### Shape Errors
248+
- Use `Shape.set_dim` (or `Shape.set_equal`) to add constraints when inference needs hints
249+
- Remember that `..var..` row variables can match zero or more axes
250+
- Check if you're unnecessarily capturing variables in einsum
251+
252+
### Type Errors with Inline Definitions
253+
- `{ x }` creates learnable parameters, not constants
254+
- Inline definitions are lifted to the unit parameter `()` scope
255+
- Sub-modules don't auto-lift - bind them before use
256+
257+
### Performance
258+
- Virtual tensors (like `0.5 + 0.5`) are inlined during optimization
259+
- Row variables allow operations to work on grouped/multi-channel data
260+
- Input axes (→) in kernels end up rightmost for better memory locality
261+
262+
## Random Number Generation Details
263+
264+
### Initialization Functions
265+
266+
OCANNL's random initialization has some important nuances:
267+
268+
1. **Default initialization is configurable** - There is a global reference that defaults to the `uniform` operation but can be changed to any nullary operation.
269+
270+
2. **Divisibility requirements** - Functions like `uniform` require the total number of elements to be divisible by certain values (they work with `uint4x32` for efficiency):
271+
- `uniform()` - requires specific size divisibility for efficient bit usage
272+
- `uniform1()` - works pointwise on `uint4x32` arrays, allows any size but wastes random bits
273+
274+
3. **Deterministic PRNG** - OCANNL uses counter-based pseudo-random generation:
275+
- Each `uniform()` call combines global seed with a unique tensor identifier
276+
- Different calls generate different streams, but deterministically
277+
- For training randomness (e.g., dropout), use `uniform_at` with `~train_step` to split the randomness key
278+
279+
**Example:**
280+
```ocaml
281+
(* Parameter init - happens once, deterministic is fine *)
282+
{ w = uniform () }
283+
284+
(* Training randomness - needs train_step for proper key splitting *)
285+
dropout ~rate:0.5 () ~train_step x
286+
(* internally uses: uniform_at !@train_step *)
287+
```
288+
289+
## Further Resources
290+
291+
- [Shape Inference Documentation](../lib/shape.mli) - Detailed einsum notation spec
292+
- [Syntax Extensions Guide](../lib/syntax_extensions.md) - `%op` and `%cd` details
293+
- [Neural Network Blocks](../lib/nn_blocks.ml) - Example implementations
294+
- [GitHub Discussions](https://github.com/ahrefs/ocannl/discussions) - Community Q&A

‎lib/nn_blocks.ml‎

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,21 @@
1-
(** This file contains basic building blocks for neural networks, with limited functionality. Feel
2-
free to copy-paste and modify as needed.
3-
4-
We follow "the principle of least commitment": where possible, we use row variables to remain
5-
agnostic to the number of axes. This flexibility often remains unused, but it makes explicit the
6-
architectural structure.
1+
(** {1 Neural Network Building Blocks}
72
8-
The einsum specifications in this file often use the single-char mode (no commas), where the
9-
spaces are entirely ignored / optional, but are used copiously for readability. *)
3+
This file contains basic building blocks for neural networks, with limited functionality. Feel
4+
free to copy-paste and modify as needed.
5+
6+
Design principles, OCANNL fundamentals, and common patterns:
7+
- "Principle of least commitment": use row variables where axis count doesn't matter
8+
- Einsum specs here often use single-char mode (no commas) but with spaces for readability
9+
- Pooling uses constant kernels (0.5 + 0.5) to propagate window dimensions
10+
- conv2d uses convolution syntax: "stride*out+kernel," (often in multi-char mode)
11+
- Input axes (before →) for kernels show intent (and end up rightmost for memory locality)
12+
- Inline params { } are always learnable and are lifted to unit parameter ()
13+
- Introduce inputs to a block after sub-block construction
14+
(sub-blocks have no automatic lifting like there is for inline definitions of params)
15+
- Always use literal strings with einsum operators when capturing variables
16+
- Avoid unnecessary variable captures in einsum operators, be mindful they can shadow
17+
other identifiers
18+
*)
1019

1120
open! Base
1221
open Operation.DSL_modules
@@ -198,9 +207,9 @@ let%op max_pool2d ?(stride = 2) ?(window_size = 2) () x =
198207
Shape.set_dim wh window_size;
199208
Shape.set_dim ww window_size;
200209
(* NOTE: projections inference runs per-assignment in a distinct phase from shape inference, so
201-
for it to know about the window size, we use a constant kernel = 1 to propagate the shape.
202-
We use a trick to create a shape-inferred constant tensor, equivalently we could write
203-
"NTDSL.term ~fetch_op:(Constant 1.) ()" but that's less concise. See:
210+
for it to know about the window size, we use a constant kernel = 1 to propagate the shape. We
211+
use a trick to create a shape-inferred constant tensor, equivalently we could write "NTDSL.term
212+
~fetch_op:(Constant 1.) ()" but that's less concise. See:
204213
https://github.com/ahrefs/ocannl/discussions/381 *)
205214
x
206215
@^+ "... | stride*oh+wh, stride*ow+ww, ..c..; wh, ww => ... | oh, ow, ..c.." [ "wh"; "ww" ]

0 commit comments

Comments
 (0)