Syntax extensions %cd and %op

Preliminaries

A self-contained dimension-capture example (compiled by the documentation check):

```ocaml doc-check=dimension_capture open Ocannl.Nn_blocks.DSL_modules

let%op sum_last x = x ++ “ab => a” [“b”]



OCANNL, and arrayjit specifically, is built around a fixed number of numeric operations, declared in `arrayjit/ops.ml`. We assign lexical operators to the binary operations, inventing novel operators if needed. For example, Rectified Linear Unit `Relu` operation, which computes `f(x) = max(0,x)`, is called `relu`, while the ReLU-Gate `Relu_gate` operation, which computes `f(x,y) = if x > 0.0 then y else 0.0`, gets the operator `-?/` in addition to name `relu_gate`. These built-in numeric operations are used to construct assignments (`Assignments.t` packaged as `Assignments.comp`). The syntax `%cd` is needed to build assignments concisely, and the assignment operators always start with `=` (unlike in C where they end with `=`). On the other hand, while the syntax `%op` helps build tensors (`Tensor.t`), they can be expressed concisely in pure OCaml. Unlike for assignments, the building blocks for tensor expressions are easy to extend. The meaningful basic ones are provided in `tensor/operation.ml`.

In OCANNL, we call a tensor that is prohibited from propagating gradients, does not have a gradient node nor backprop code, a _non-differentiable tensor_. Accordingly we can call the "plain" tensors with a gradient node _differentiable tensors_. Expressions in the `%cd` syntax will sometimes build new non-differentiable tensors as components of assignments (they will never build new differentiable tensors). The syntax extensions make the following assumption:

- `%cd` assumes that any extension point will be in the scope of a module `NTDSL` that provides at least the functionality of `Operation.NTDSL`.
- `%op` assumes that any extension point will be in the scope of a module `TDSL` that provides at least the functionality of `Operation.TDSL`.
- Both extensions assume `Tensor` (from the `Ocannl` wrapper) is in scope.

Functions inside `Operation.NTDSL` use `~grad_spec:Prohibit_grad` when calling into `Tensor`, making the resulting tensors non-differentiable. Functions inside `Operation.TDSL` use `~grad_spec:If_needed`, which will make the tensors non-differentiable when the gradient is not needed -- except for `TDSL.param`, which makes the final parameter tensor differentiable after evaluating its initializer as forward-only code. A `Prohibit_grad` operation cuts the gradient at that operation whatever its operands are: `Tensor.op` builds no gradient node and composes no backpropagation code for it, so applying an `NTDSL` operation to a differentiable operand -- `NTDSL.O.(x @^^ spec)` inside a `%op` body, as `Nn_blocks.softmax` does for its row max -- detaches the result without a copy. `Operation.stop_gradient` is the identity instance of the same mechanism.

The extension points open `NTDSL.O`, resp. `TDSL.O`, for the scope of the extension point, to expose the corresponding operators.

### The %oc anti-quotation and the unit-parameter heuristic

Within `%op` and `%cd` contexts, expressions typically undergo transformation to build tensors or assignments. However, OCANNL uses two mechanisms to preserve pure OCaml expressions:

#### Unit-parameter heuristic (automatic in %op)

In the `%op` syntax, when a function application contains a unit `()` argument, all arguments appearing **before** the unit are automatically preserved as pure OCaml expressions. This aligns with OCANNL's design pattern where configuration happens before the unit parameter:

```ocaml
(* Arguments before () are automatically preserved as OCaml *)
let%op my_fn ~label x = 
  other_fn ~label:(("prefix_" ^ name) :: label) ~config:value () x
  (* label and config are preserved; x after () is transformed *)

Explicit %oc anti-quotation

For cases where you need explicit control or the heuristic doesn’t apply, the %oc (mnemonic: “OCaml”) anti-quotation escapes from the transformation context:

(* Force preservation even after () or in edge cases *)
let%op special = process_data data [%oc complex_ocaml_expr]

The %oc extension expects a single expression and returns it unchanged. Use cases: - Overriding the unit-parameter heuristic when needed - Preserving expressions in contexts without a unit parameter - Escaping from the DSL in %cd contexts (which don’t use the unit heuristic)

Primitive operations

To accommodate stylistic preferences, OCANNL supports both curried and uncurried syntaxes for primitive operation application. Binary operators have infix and alphabetic identifiers. The infix forms stem from the following restriction: in the %cd syntax, the assignment is always an infix operator, and it needs to pick the accumulation operation. In %op, alphabetic names are available for the binary operations its O module provides: add, sub, mul, div, lt, le, eq, and ne (mul and div are pointwise). The remaining binary primitive names in the table below are %cd-only; in %op, those identifiers retain their ordinary OCaml meaning.

In %op, the uncurried form op (a, b) or op (a, b, c) of a binary or ternary primitive operation is the curried call op a b or op a b c: where (c, x, y) is where c x y, and ( + ) (x, y) is x + y. The operands are %op expressions, so they can hold inline declarations. The identifier must be one the DSL’s O module binds, as for the curried form.

The unary primitive operations:

Identifier Default projection Constructor in Ir.Ops
id pointwise Identity
relu pointwise Relu
sat01 pointwise Satur01
exp pointwise Exp
log pointwise Log
exp2 pointwise Exp2
log2 pointwise Log2
sin pointwise Sin
cos pointwise Cos
sqrt pointwise Sqrt
recip pointwise Recip
recip_sqrt pointwise Recip_sqrt
neg pointwise Neg
tanh pointwise Tanh_approx
not pointwise Not
uint4x32_to_prec_uniform dedicated Uint4x32_to_prec_uniform

The binary primitive operations:

Identifier Infix operator Default projection Constructor in Ir.Ops Assignments
fst -@> pointwise Arg1 none
snd -/> pointwise Arg2 =:
add + pointwise Add =+, =:+
sub - pointwise Sub =-, =:-
mul * none Mul =*, =:*
div / none Div =/, =:/
pow ** pointwise ToPowOf =**, =:**
relu_gate -?/ pointwise Relu_gate =?/, =:?/
sat01_gate -?^ pointwise Satur01_gate =?^, =:?^
lt < pointwise Cmplt none
le <= pointwise Cmple none
eq = pointwise Cmpeq none
ne <> pointwise Cmpne none
or_ \|\| pointwise Or =\|\|, =:\|\|
and_ && pointwise And =&&, =:&&
mod_ % pointwise Mod none
max @^ pointwise Max =@^, =:@^
min @- pointwise Min =@-, =:@-
threefry4x32 ^^^^ pointwise Threefry4x32 =^^^^, =:^^^^

The ternary primitive operations:

Identifier Default projection Constructor in Ir.Ops
where pointwise Where
fma compose-accumulate FMA

The interpretation functions also state the semantics:

let interpret_unop op v =
  let open Float in
  match op with
  | Identity -> v
  | Relu when v >= 0. -> v
  | Relu -> 0.
  | Satur01 when v <= 0. -> 0.
  | Satur01 when v >= 1. -> 1.
  | Satur01 -> v
  | Exp -> exp v
  | Log -> log v
  | Exp2 -> 2. ** v
  | Log2 -> log v / log 2.
  | Sin -> sin v
  | Cos -> cos v
  | Sqrt -> sqrt v
  | Recip -> 1. / v
  | Recip_sqrt -> 1. / sqrt v
  | Neg -> ~-.v
  | Tanh_approx -> tanh v
  | Not -> if v = 0. then 1. else 0.
  | Uint4x32_to_prec_uniform -> failwith "NOT IMPLEMENTED"

let interpret_binop op v1 v2 =
  let integer_power base exponent =
    let open Float in
    let rec loop base magnitude result =
      if magnitude = 0. then result
      else
        let result = if Stdlib.Float.rem magnitude 2. = 1. then result * base else result in
        let magnitude = Stdlib.Float.floor (magnitude / 2.) in
        if magnitude = 0. then result else loop (base * base) magnitude result
    in
    let result = loop base (Float.abs exponent) 1. in
    if exponent < 0. then 1. / result else result
  in
  let open Float in
  match op with
  | Arg1 -> v1
  | Arg2 -> v2
  | Add -> v1 + v2
  | Sub -> v1 - v2
  | Mul -> v1 * v2
  | Div -> v1 / v2
  | ToPowOf when is_finite v2 && is_integer v2 -> integer_power v1 v2
  | ToPowOf -> v1 ** v2
  | Relu_gate -> if v1 > 0.0 then v2 else 0.0
  | Satur01_gate -> if v1 > 0.0 && v1 < 1.0 then v2 else 0.0
  | Max -> max v1 v2
  | Min -> min v1 v2
  | Mod -> v1 % v2
  | Cmplt -> if v1 < v2 then 1. else 0.
  | Cmple -> if v1 <= v2 then 1. else 0.
  | Cmpeq -> if v1 = v2 then 1. else 0.
  | Cmpne -> if v1 <> v2 then 1. else 0.
  | Or -> if v1 <> 0. || v2 <> 0. then 1. else 0.
  | And -> if v1 <> 0. && v2 <> 0. then 1. else 0.
  | Threefry4x32 -> ...

let interpret_ternop op v1 v2 v3 =
  let open Float in
  match op with Where -> if v1 <> 0. then v2 else v3 | FMA -> (v1 * v2) + v3

Known finite integral powers use multiplication and squaring, followed by a reciprocal for negative exponents. The exponent retains its host-float value without conversion to a machine int. Exponent zero returns 1 even for zero or NaN bases; nonzero NaN powers propagate NaN, and odd zero and infinity powers preserve their sign. Generated arithmetic rounds at f32 for single and narrow storage, and at f64 for double storage; this multiplication policy includes its overflow and underflow behavior. Host interpretation uses f64, so constant integer powers remain in the IR for target-precision code generation, including bases that simplify to constants. Fractional and dynamic exponents use the backend’s floating-power operation. Optional unrolling applies only to small nonnegative exponents and does not duplicate scoped computations.

The syntax for %op

The %op syntax is simpler than the %cd syntax since it relies more on regular OCaml expressions. For example, we can write without syntax extensions:

  let hid_dim = 8 in
  let w = TDSL.param "w" in
  let b = TDSL.param ~output_dims:[ hid_dim ] "b" in
  let layer x = TDSL.O.( relu(w * x + b) ) in
  ...

Since TDSL.O is opened for the scope of an extension point %op:

  let hid_dim = 8 in
  let w = TDSL.param "w" in
  let b = TDSL.param ~output_dims:[ hid_dim ] "b" in
  let%op layer x = relu(w * x + b) in
  ...

Using inline declarations, this becomes more concise:

  let hid_dim = 8 in
  let%op mlp_layer x = relu({ w } * x + { b; o = [ hid_dim ] }) in
  ...

When there is a function directly under the %op extension point, like in the example above, or directly under a function taking a unit parameter (), the function parameter (to the right of ()) should be a tensor. That’s because automatic shape schemes take such an unlabelled parameter to be a tensor (see Automatic shape schemes), and because, in a function without (), %op uses this tensor’s (value’s) label to enrich the label of the resulting tensor (see Label from function argument). A labelled or optional parameter may be anything, with or without (): without (), a labelled one is read as a tensor only on evidence that it is one (see Label from function argument).

When the declaration is followed by a literal float, the float provides the initial value to initialize the tensor. Otherwise, the tensor value cells are initialized with the default parameter initialization: a centered, scaled packed uniform () over [-0.25, 0.25) (Operation.centered_uniform_param_init with scale = 0.5), configurable by setting the reference TDSL.default_param_init.

The syntax for %cd

The basic building blocks of the %cd syntax are individual assignments, separated by semicolons. The assignments, represented via Assignments.Accum_binop and Assignments.Accum_unop, are in full generality accumulating:

type Assignments.t =
   ...
  | Accum_binop of {
      initialize_neutral : bool;
      accum : Ops.binop;
      op : Ops.binop;
      lhs : Tnode.t;
      rhs1 : buffer;
      rhs2 : buffer;
      projections : Indexing.projections Lazy.t;
    }
  | Accum_unop of {
      initialize_neutral : bool;
      accum : Ops.binop;
      op : Ops.unop;
      lhs : Tnode.t;
      rhs : buffer;
      projections : Indexing.projections Lazy.t;
    }

For example the binary case in pseudocode: if initialize_neutral then lhs = 0; lhs = lhs accum (rhs1 op rhs2) (assuming the neutral element of accum is 0). The representation also has a field projections which determines which loops should be run and how the tensor nodes should be indexed to perform the computation.

The basic %cd syntax for assignments has the form: <lhs> <asgn-op> <primitive-op-application[rhs1, rhs2?, rhs3?]>. See Primitive operations for the syntax of primitive operation application, where <rhs1>, <rhs2> (for binary and ternary ops), <rhs3> (for ternary ops) are subexpressions. <asgn-op> starts with =, followed by : only if initialize_neutral is true, then followed by the operator syntax variant of a binary primitive operation. The fields <lhs>, <rhs1>, <rhs2>, <rhs3> will often be either special-purpose identifiers (specifically v, t, t1, t2, t3, g, g1, g2, g3) or identifiers bound to tensors. <rhs1>, <rsh2>, <rsh3> will also often be (non-differentiable) tensor expressions. The notation <tensor>.grad stands for the gradient node of the given tensor. For more about “slot fillers”, and to learn about the operators +* and ++, see the section further features of the syntax extension %cd.

How is the projections field determined? projections can be given explicitly as a labeled argument ~projections. If they aren’t but %cd realizes there is a ~projections parameter in scope, it uses it – see tensor/operation.ml where this option is used to define tensor operations. If instead of ~projections a ~logic labeled argument is given, the string passed is used to determine projections. ~logic:"." means a pointwise operation. ~logic:"@" means an “output axes of rhs2 match input axes of rhs1” operation (matrix multiplication is a special case). ~logic:"T" means transpose of input and output axes. The string passed to ~logic can also use OCANNL’s generalization of the einsum notation, allowing arbitrary permutations and reductions of axes. If no information is given, the default depends on the primitive operation, but it is almost always a pointwise operation.

Restriction: Using ~logic:"@" with / (div) or ** (pow) is a compile-time error: those combinations read like matrix inverse or matrix power but compute neither (e.g. sum_k a[i,k] / b[k,j]). Use ~logic:"." for pointwise division/power, or an explicit einsum string for a custom contraction. Other binary ops with ~logic:"@" (including *, +, max, etc.) remain accepted — semiring-style contraction over non-multiply ops is intentional (cf. the tropical @^+ operator).

Here we see an example of tensor multiplication – extending matrix multiplication to arbitrary number of axes – multiplying a by b to get c. In =:+, = is required to separate the assigned-to part from the computation, : clears-out c before the computation, + selects addition to accumulate the results.

c =:+ a * b ~logic:"@"

Compare the following two ways of updating a parameter p:

p =+ learning_rate * p.grad ~logic:"."

and:

p =+ learning_rate *. p.grad

In the first case, we have a binary assignment calculated pointwise. The resulting representation is Accum_binop where accum is Add and op is Mul (multiplication). In the second case, *. is not recognized as one of the built-in operators. This leaves the expression learning_rate *. p.grad un-transformed. Since (*.) is bound in NTDSL.O to pointwise tensor multiplication, this creates an intermediate tensor, that is then added onto p. The resulting representation is Accum_unop where accum is Add and op is Identity. Both variants end up with the same result, and even with the same computation, because the second variant’s computation will get optimized (unless configured not to).

Reading the lhs, and state that persists across runs. Only the =:-prefixed accumulating operators (=:+, =:*, =:@^, …) clear the lhs first; plain =: is an assignment that does not, and =+, =-, =*, … accumulate onto whatever the lhs holds. So an assignment whose right-hand side names its own lhs under =:, and every non-clearing accumulation, reads the lhs’s previous value: what an earlier statement of the same routine wrote, or else what the previous run of the routine left. A node read that way is read-before-write, and the compiler keeps it in a device buffer that persists between runs — never virtual, never inlined. This is how optimizer state is declared, with no extra annotation: Train.sgd_update ~momentum writes { sgd_momentum } =: (!.momentum *. sgd_momentum) + sgd_delta, and the buffer carries the velocity from one step to the next. Forcing such a node virtual (e.g. Train.set_virtual) is refused at compilation.

Advanced note: when a ~projections parameter is in scope but no assignment-specific ~projections argument is given – the typical case in tensor/operation.ml – the actual projections field for an assignment is computed by transforming the projections parameter according to hints regarding how tensor nodes relate to the given projections. Specifically, the identifiers rhs1, t1, v1, g1 are “slot RHS1” of the projections, rhs2, t2, v2, g2 are “slot RHS2”, lhs, t, v, g are “slot LHS”. Scalar constants are provided the projection directly, to make the automated derivation more expressive; this is supported both for literals, and (heuristically) for !. and !.. embedding operators.

Projection slot detection by naming convention

In addition to the special identifiers (t, t1, t2, lhs, rhs1, etc.), the %cd syntax can detect projection slots from identifier and inline tensor definition names using prefix/suffix patterns. This is essential when defining backpropagation code that needs intermediate tensors with specific projections and shapes.

The naming convention patterns:

Prefix/Suffix Detected Slot Shape Source
lhs_* or *_lhs LHS output shape (from t)
rhs_* or *_rhs RHS1 first input shape (from t1)
rhs1_* or *_rhs1 RHS1 first input shape (from t1)
rhs2_* or *_rhs2 RHS2 second input shape (from t2)
rhs3_* or *_rhs3 RHS3 third input shape (from t3)
pspace_* or *_pspace product space the operation’s projections (see below)

This applies to: - Inline tensor definitions: { cond_rhs1 } declares a tensor with slot RHS1 and shape of t1 - Identifier references: When sum_rhs1 is used in an expression, it’s recognized as having slot RHS1

The product-space slot (*_pspace): the tensor is laid out over the operation’s full product space — the result’s axes followed by the reduced-over (contracted) axes — and is indexed by the identity projection over the product iterators. This is the index space in which per-(result, reduced position) facts live. The inline-declared tensor’s shape is unified with the operation’s product-space proxy shape (Shape.product_space_shape): the result’s rows with the contracted axes appended (fixed-index axes pinned to dimension 1 — their projections are pinned, so they are never product axes). At most one whole row variable may be reduced over (it becomes the row variable of the appended flank); otherwise, or for a non-einsum operation, a shape-inference error is raised. The identity projection (Ir.Indexing.prod_project_for) skips dimension-1 axes and pairs the rest with product components by extent (first-fit, so the layout order may deviate from the product order); a leftover on either side raises at lowering.

Why this matters for gradient computation: In operations like max pooling or tropical convolution, the gradient must flow back to exactly the (window, position) pairs that achieved the argmax. A gate materialized over an operand’s index space is lossy whenever reduction windows overlap (stride < window): several windows write conflicting bits through a non-accumulating assignment and the last write wins, silently corrupting gradients (gh-ocannl-512). The product-space gate has one cell per (result position, reduced position) pair, where the equality check is single-writer and therefore exact.

Example from the tropical operation’s gradient computation:

let%cd grad_asn ~t ~g ~t1 ~t2 ~projections =
  (* Product-space intermediates: one cell per (output, kernel-position) pair. *)
  { sum_pspace } =:@^ add (t1, t2);      (* the candidate value of each pair *)
  { cond_pspace } =: eq (t, sum_pspace); (* true where the pair achieved the argmax *)
  { valid_pspace } =:|| eq (t1, t1);     (* false only for clamped-out (padded) accesses *)
  cond_pspace =&& valid_pspace;
  g1 =+ where cond_pspace g 0;           (* gradient flows to argmax input positions *)
  g2 =+ where cond_pspace g 0            (* gradient flows to argmax kernel positions *)
in

(The initializing max-accumulation =:@^ gives pairs whose t1 access is clamped out of range by a padded window spec the max-neutral -inf; the =:|| validity mask handles the corner where t itself is -inf — an out-of-range read substitutes the statement’s accumulation neutral, which for Or is 0 where an in-range read contributes 1.)

Important: The naming convention affects both projection slot assignment and shape inference. In addition to determining which projection from ~projections to use when indexing the tensor, a shape equality constraint is generated between the inline-defined tensor and the corresponding operation tensor assumed to be in scope: t for *_lhs, t1 for *_rhs and *_rhs1, t2 for *_rhs2, etc. This means the shape from the tensor’s initialization is unified with the shape of the operation component. For *_pspace the unification target is the operation’s product-space proxy shape instead, and the declared tensor should not be given explicit dimensions.

Numeric and N-dimensional array literals

Both %cd and %op extensions use a shared syntax for N-dimensional array literals. %cd uses NTDSL.number and NTDSL.ndarray functions, while %op uses TDSL.number and TDSL.ndarray functions. (This is just for consistency: TDSL.ndarray invokes Tensor.ndarray ~grad_spec:If_needed, which will figure out the gradient is not needed and will make the tensor non-differentiable.)

Numbers are a special case: an array of (output) dimension 1.

N-dimensional array literals combine the list, tuple and array syntaxes to strictly distinguish between output, input and batch axes:

For example, [ (1, 2, 3); (4, 5, 6) ] is a mathematical matrix converting 3D vectors into 2D vectors.

OCANNL supports dimension bases (semantic annotations on dimension values, a multi-character tag such as rgb or heads). You label an axis by writing an OCaml type-annotation on the axis container inside a tensor literal — the annotation names that axis’s basis:

Annotations compose with nesting: in [ ((1., 2.) : feat) ] the inner tuple labels the input axis feat while the enclosing output list is unannotated. Only a bare type-constructor name is read as the basis (qualified, parameterised, tuple, arrow and variant types are ignored). These bases then propagate to specify bases of other dimensions in other tensors, via shape inference.

Outside literal axis-basis syntax, (relu x : Tensor.t) preserves the type annotation and translates the application inside it as usual.

As a convenience, the same annotation on a bare numeric literal labels the size-1 output axis of the resulting scalar: (2.0 : q) is shorthand for ([ 2.0 ] : q).

(This replaces an older 'q' 2.0 char-literal form, which could not carry multi-character tags.) Unannotated axes carry the reserved default basis (an atom that does not fuse with a named basis, so an unlabelled axis will not silently merge with a labelled one), and the reserved bcast_if_1 basis is the claim-free broadcast top that a size-1 axis stretches from — write it explicitly (e.g. ([ 1. ] : bcast_if_1)) when you want a size-1 axis to broadcast. Example: let%op y = ({ hey } * ([ 2.0 ] : q)) + ([ 1.0 ] : p) in ....

Wildcard bindings

When an extension is over a wildcard (ignore result) binding: let%cd _ = ... and let%op _ = ..., the generated code is wrapped in Tensor.with_unchanged_roots, to prevent it from upsetting rootness checks. The use-case for writing %op and %cd notations with ignored result is to generate additional shape inference constraints.

Inline declarations

Both %cd and %op syntaxes support inline declarations of tensors. For %op these are differentiable, for %cd non-differentiable tensors.

A declaration site uses the record syntax. The key difference between the two extensions:

Both syntaxes support additional record fields that map directly to labeled arguments of the tensor creation functions (see Tensor module signatures):

Note: for %op inline parameter declarations, initialization expressions are forward-only. If the root operation is not qualified with a module name, it is qualified with NTDSL; TDSL.param then makes only the final parameter tensor differentiable. There are also special cases for literal constants so the resulting tensor is initialized with those constants and the final parameter remains differentiable.

Examples:

The tensor name is bound to the newly created tensor, and the record expression itself evaluates to the tensor. The scope of the binding is the full scope of the extension point, even if the declaring record appeared in the body of a function that’s inside the extension point scope (except for %op there is a special case of functions taking a unit parameter () discussed below – inline definitions are introduced once () is applied). The first element of the label of the created tensor is the name that introduced it.

For %cd, inline declarations are allowed both in the assigned-to position (left-hand side) of assignments and in standalone tensor expressions. When used in assignments, one of the tensors on the right-hand-side is picked to provide additional label information if possible. In particular, tensors that are function parameters inside the scope of the extension point, cannot be picked to provide label information, as they would escape their scope at the point the tensor is created. Inline declarations are still prohibited within the right-hand side of assignments to discourage over-use in locations with less label information. Example showing two tensor nodes declared inline, both of them include the label of the param p in their labels:

let sgd_one ~learning_rate ?(momentum = 0.0) ?(weight_decay = 0.0) ?(nesterov = false) p =
  [%cd
    { sgd_delta } =: p.grad + (!.weight_decay *. p);
    if Float.(momentum > 0.0) then (
      { sgd_momentum } =: (!.momentum *. sgd_momentum) + sgd_delta;
      if nesterov then sgd_delta =+ !.momentum *. sgd_momentum else sgd_delta =: sgd_momentum);
    p =- learning_rate *. sgd_delta]

Because the scope of an inline declaration is the whole extension point, %cd hoists inline declarations above an enclosing match or if: the Pexp_match arm (via handle_cases) and the if ... then ... else arm of tensor/ppx_cd.ml both collect their branches’ bindings with reduce_vbss and lift them out of the branch. So an arm that declares { x } and an arm that only reads x name the same tensor, and branches need not be duplicated just to share one.

The declaration may appear in only one arm, though: a second { x } for the same name is rejected with ppx_ocannl %cd: duplicate inline declaration of no-gradient tensor x (raised by assignment in tensor/ppx_cd.ml). Declare the tensor in whichever arm reads most naturally – typically the one whose right-hand side carries the best label information – and refer to it by bare name in the others.

Train.sgd_one in lib/train.ml is the worked example: its %cd payload is written once, with { sgd_delta } and { sgd_momentum } each declared in a single arm of the grad_scale / update_gate matches and read by name everywhere else, so the arms differ only where the gating actually changes the emitted assignment. (The momentum block was duplicated across the gate branches until gh-ocannl-772, on the mistaken assumption that hoisting past a match did not work.)

Inline declarations can also be used outside of assignments for creating non-differentiable tensors, to mimic the behavior of %op but without the burden of initialization that a parameter would introduce:

  let%cd mlp_result = mlp { point } in
  let _, result_routine =
    Train.to_routine sgd_routine.Context.context IDX.empty
      [%cd ~~("mlp infer"; mlp_result.forward)]
  in
  let callback (x, y) =
    (* Context-mediated value access (gh-ocannl-333): write the input on-device, then read the
       result through the context with the [At] accessors taking a [(context, tensor)] pair. *)
    ignore (Context.set_values ctx point [| x; y |] : Context.t);
    Train.run ctx result_routine;
    Float.((ctx, mlp_result).@[0] >= 0.)
  in

For %op, the declaration is allowed anywhere. If there is a unit () parameter in the function, the scope of inline-declared tensors is delimited at that parameter. The tensors are defined right after the unit parameter. If there is a labeled or optional parameter with label label before the unit parameter (~label, ?label, ?(label = [ "net" ])), the inline-declared tensors will use that parameter (which should be of type string list; an omitted ?label adds nothing) to enrich their labels. Example showing two param tensors declared inline, with scope delimited by () and labels enriched by the label parameter:

let%op mlp_layer ~label ~hid_dim () x = relu ({ w } * x + { b; o = [ hid_dim ] })

Implementation strategy for the initialization syntax

To maintain the familiar concise syntax, yet allow for configurability during initialization, the %op syntax substitutes the operator function applied at the root of the initialization expression by prefixing the function identifier with NTDSL. Only unqualified identifiers get prefixed, and %oc is an escape hatch to prevent prefixing even for unqualified identifiers.

Using OCANNL’s generalized einsum notation

For einops users: If you’re familiar with einops, see einops_comparison.md for a side-by-side mapping of einops operations to OCANNL’s notation.

As we mentioned above, in the %cd syntax you can set up an arbitrary assignment with projections derived from a generalized einsum specification, by passing the specification as a string with the ~logic label. However, both the %cd and %op syntaxes support built-in operators that take an einsum specification: +* binding to NTDSL.einsum resp. TDSL.einsum, and ++ binding to NTDSL.einsum1 resp. TDSL.einsum1. +* is a “ternary” operator, binary wrt. tensor arguments, and ++ is a binary operator, unary postfix wrt. tensor arguments. There are even more einsum operators: binary @^+ binding to tropical and +++ binding to outer_sum; unary @^^ binding to einmax1 (the operator-to-function tables are einsum_binary_ops/einsum_unary_ops in tensor/ppx_shared.ml). The ternary einsum has no operator: write the named form einsum3 "ij; jk; km => im" a b c (outside the syntax extensions, TDSL.einsum3 spec a b c () or NTDSL.einsum3 spec a b c (), as for the other named forms). When the einsum specification is a literal string, we support two syntax patterns: the string can either directly follow the operator (infix-style notation), or the string can follow the second argument (mixfix-style notation). When the spec string is an identifier, it must directly follow the operator.

+*, +++ and ++ use addition for the accumulation operation; @^+ and @^^ use maximum. The pairwise operation combined with the accumulation: +* multiplies (the classic einsum contraction), +++ and @^+ add (so @^+ is max-reduce with add — the tropical semiring), and the unary ++/@^^ reduce their argument directly (add-reduce resp. max-reduce). You can verify that looking at the definitions of Operation.einsum, Operation.einsum1, etc. You can find examples of +* and ++ behavior in the test suite einsum_trivia.ml and in nn_blocks.ml. A frequent use-case for ++ is to sum out all axes of a tensor:

  let%op scalar_loss = (margin_loss ++ "...|... => |->0") /. !..batch_size in
  ...

where (!..) converts an integer into a constant tensor. Note the |->0 result spec: the bare separators close the result’s batch and input rows, so the batch axes are reduced rather than broadcast through (see below).

Syntax of the generalized einsum notation

The specification syntax has two modes:

A literal operator capture list is checked by the ppx: every captured name must occur as a label or a named row variable in the spec. For example, "ab => a" ["ab"] is refused at the capture literal: that spec names axes a and b. To name one axis ab, use multichar mode, e.g. "ab, => ab," ["ab"]. Writing "a, b => a" still names two separate axes.

Whitespace never separates axes. In single-char mode it is ignored (s a is the two axes sa); in multichar mode axes need commas between them, so s a; s b => s a^b is a parse error (the ^ selected multichar mode) whose message names the fix, s, a; s, b => s, a^b.

The syntax of a generalized einsum spec has two variants:

Recall that a tensor shape is composed of three rows, i.e. sequences of axes: batch, input and output axes. Correspondingly, a shape spec in the notation can be:

A row whose kind separator is omitted altogether reads as the context ellipsis: x is equivalent to ...|...->x, so batch and input axes broadcast through terse specs by default. To require an empty row, write the separator with an empty row spec: | ->x means no batch and no input axes, |x no batch axes, ->x no input axes.

The notation for a row is composed of sequences of row specs, and an optional row variable spec. A row variable tracks broadcasting. The syntax of a row:

The syntax of a row variable:

The syntax of an axis spec:

Examples:

Affine indexing for convolutions and pooling

The affine axis syntax enables convolution and pooling operations directly in einsum notation. The semantics:

Important constraint for valid convolution: The formula must hold exactly, meaning (input_size - effective_kernel_span) must be divisible by stride. Otherwise, shape inference will fail with “incompatible stride” error.

General rule (use_padding = false case): (input_size - effective_kernel_span) mod stride = 0

For example, with stride=2, kernel_size=2, dilation=1: - effective_kernel_span = 1 + (2-1) * 1 = 2 - A 4x4 input gives output_size: 4 = 2 * (output - 1) + 2 → output = 2 ✓ - A 5x5 input would fail: (5 - 2) mod 2 = 1 ≠ 0 → shape inference error

With stride=2, kernel_size=3, dilation=1: - effective_kernel_span = 1 + (3-1) * 1 = 3 - A 9x9 input works: (9 - 3) mod 2 = 0 → output = 4 ✓ - A 10x10 input fails: (10 - 3) mod 2 = 1 ≠ 0 → shape inference error

Examples:

Axis concatenation with the ^ operator

The ^ operator in einsum specifications creates a concatenated axis from multiple components. This enables:

The syntax a^b (or a^b^c etc.) creates a single axis of iteration that first iterates over component a, then over component b, etc. The components are axis labels. Since ^ selects multichar mode, they are identifiers and the spec’s axes are comma-separated: sa; sb => sa^b reads as the single axes sa, sb and sa^b; concatenating a and b under a shared leading axis s is written s, a; s, b => s, a^b.

Examples of concatenation patterns (using vector notation for simplicity):

Pattern Description
a; b => a^b Concatenate vectors: result contains all of a then all of b
a^b => a Extract prefix: take the first part of a vector
a^b => b Extract suffix: take the last part of a vector
a => a^b Replace prefix: assign to first part, leaving suffix unchanged
b => a^b Replace suffix: assign to last part, leaving prefix unchanged
a^b^c => b Extract middle: requires knowing sizes of a and c
b => a^b^c Replace middle: requires knowing sizes of a and c
a^b^c => a^c Drop a segment: the result is a then c
a^c => a^b^c Write around a segment: b is left unwritten (zero in a fresh result)
a^b^c => c^b^a, a^b^c => c^a Reorder segments, dropping or not

Integer constants in concatenation: an integer component specifies the size of that segment, rather than indexing into a fixed dimension. For example: - 3^a => a: skip 3 elements at the beginning of the input - a => a^3: assign to all but the last 3 elements of the result - 3^a^5 => a: extract the middle portion, skipping 3 at the start and 5 at the end

No other slot names an integer component, so it is skipped: an extraction does not read it, a replacement does not write it. An integer component must be positive.

Integer components carry the default basis (a short-term rule, gh-ocannl-2093): an integer is a size with no basis of its own, so a concatenation of integers alone sums to an unannotated axis, which does not fuse with a named one. On an operand x whose axis is rgb of size 5, x ++ "2^a => a" is accepted, the variable a taking the rgb claim (its binding records it, so it cannot take a second named basis later); x ++ "2^3 => " is refused with solved dimensions for axis: different bases. A fixed claim-free tag for integers would be unsound: a solved size that meets another records no basis, so a claim-free 5 could meet rgb 5 and later hsv 5, merging the two silently. Workaround: name a segment, as in 2^a. The principled answer, an integer component taking a slice of its partner’s basis, is the algebra of concatenation bases of gh-ocannl-2095. Pinned in test/einsum/concat_integer_components.ml and test/operations/test_concat_dim_solver.ml.

Sizing the components: every component must have a size. A component another slot names gets it from that tensor; the others take an integer, or a captured variable set with Shape.set_dim (x ++ "p^o^q => o" ["p"; "q"], then Shape.set_dim p 2 and Shape.set_dim q 1). An extraction’s result closes down to its operand, so it cannot size o from a use site: in p^o => o either p or the result is sized explicitly, and a middle form needs two of its three components. Shape inference never guesses an unsized component to 1, as it would close an unconstrained axis: it refuses with a message naming the component. (A block spec’s component no operand supplies may still be inferred empty.)

Extraction and replacement: a concatenated axis exactly one of whose components the other side names is a slice: its index is the named component’s index shifted by the sizes of the components before it, and the other components are iterated by nothing. ++ and ++^ results are zero-initialized, so a replacement into a fresh result zero-pads it (a => 2^a^1); a %cd assignment target =: source ~logic:"a => a^b" writes the segment and leaves the rest of target unchanged. The gradient of an extraction flows into the selected segment only. These forms are executed against oracles in test/einsum/concat_extraction.ml.

Segments correspond by name: when two slots both concatenate components of one axis, each component is matched by its label, and each slot places it at its own offset. A segment one slot does not list is one that slot takes no part in, so the two sides may differ in order and in what they list: a^b^c => c^b^a reverses the segments, a^b^c => a^c drops b, a^b^c => c^a does both, and a^c => a^b^c writes a and c around a b no operand fills – zero under ++, its prior contents under %cd =:. An integer component is matched by nothing, as everywhere: a^2^c => a^c drops two cells, a^c => a^2^c leaves two unwritten. Gradients follow the segments back, and a dropped segment gets none. These forms, binary ones (a^b^c; a^c => a^c) among them, are executed in test/einsum/concat_segment_correspondence.ml (gh-ocannl-2087).

What a concatenated axis cannot say: a concatenated axis is one axis of alternatives, so a position lies in one segment. Five forms therefore have no meaning, and are refused with a shape error when the operation is built, whatever the sizes (gh-ocannl-2017):

The same forms on an operand are defined where every cell has a reading: an operand concatenation no component of which the result names is a reduction over all its cells (2^3 =>, p^q => c), and every other access to that operand sees a plain axis of the summed size. When compiling, an operand that no segment reaches is refused too: in a; b => a^b, b the result’s plain b rules out the segment a that the first operand fills, and in a^b, b => a the result and the operand need different segments. The forms are executed in test/einsum/concat_axis_consumers.ml.

N-ary einsum for multiple tensors: The einsum parser supports any number of RHS tensors separated by semicolons. This enables operations like:

(* Concatenate 3 output-axes-only tensors along the first axis *)
(a, b, c) ++^ "x, ...; y, ...; z, ... => x^y^z, ..."

Syntax for %op vs %cd: In the %op extension, use the ++^ operator for concatenation and block specs. For %cd, single-argument concatenation (slicing, partial updates) works with the existing ~logic:"..." syntax:

(* Assign to prefix of target, leaving suffix unchanged *)
let%cd update_prefix ~target ~source =
  target =: source ~logic:"a => a^b"

Multi-argument syntax for %cd (needed for tensor concatenation with multiple sources) is still being designed. The natural choice [rhs1; rhs2] conflicts with the block tensor syntax (implemented in %op).

Block tensor syntax

The tensor literal syntax is generalized to support block tensor construction: when a list, tuple, or array literal inside %op contains tensor expressions (rather than numeric literals), it desugars into a call to the stack operation (which composes einsum1 unsqueeze + concat along a fresh leading axis). The result has rank = operand rank + 1 (true stacking, a new axis).

Axis mapping by syntax form:

Syntax Axis kind Example
( , ) Input (ta, tb) stacks along a new input axis (top-level %op expressions only)
[ ; ] Output [ta; tb] stacks along a new output axis
[| ; |] Batch [|ta; tb|] stacks along a new batch axis

For example, [ta; tb] desugars to stack `Output [| ta; tb |], which generates: 1. An unsqueeze einsum spec introducing a new size-1 output axis per component: ...|...->... => ...|...->0,..., applied to each operand via einsum1. 2. A concat spec joining the expanded operands along that axis: ...|...-> bt0,...; ...|...-> bt1,... => ...|...-> bt0^bt1,..., applied via concat.

Disambiguation rule: if the first leaf of the nested literal is a numeric constant (int or float, descending through type annotations like (2.0 : rgb)), the expression is treated as an ndarray constant (existing behavior). Otherwise it is a block tensor. To include scalar constants alongside tensors, wrap them with !. or !.. (e.g., [!.1.0; ta]).

Notes:

Capturing the dimensions of selected axes for further computation or to add shape constraints

The einsum operators (+*, @^+, +++, ++, @^^) and the concatenation operator ++^ accept an optional list of strings argument after the specification string. When passed, the strings should be some of the identifiers used in the specification. Both dimension variable and row variable labels are supported. This will introduce bindings for Indexing.variable_ref objects per call: at the top of the innermost function body enclosing the operation, or, outside of any function, at the same point as the inline parameter definition bindings. The ppx passes these objects with the ~capture_dims argument to einsum, einsum1, resp. concat. The bound objects can be used to add shape constraints via Shape.set_dim and Shape.set_equal, or with Operation.embed_dim or its alias Operation.TDSL.O.dim to embed the solved dimension of the corresponding variable (as a number) into a tensor expression. For a row variable, the number will be the product of the dimensions it resolved into.

%cd assignments also accept a trailing capture list on a literal ~logic specification, for unary, binary and ternary assignments. For example, let%cd comp = tgt =: id x ~logic:"a^c => a^b^c" ["a"; "b"] in binds a and b with the same per-call scope and label checks as operator captures. Captures apply to einsum specs only: the shortcuts ".", "T" and "@" reject nonempty capture lists, but accept []. Shortcut spellings take precedence over einsum parsing: "." accepts unary, binary and ternary assignments, "T" only unary, and "@" only binary or ternary; other arities are expansion errors with or without [] (use an explicit einsum spec such as "T => T" for an axis named T). Shape.set_dim a 2; Shape.set_dim b 1 sizes those concat segments before compilation. This assignment writes the a and c segments; the existing b cells of tgt retain their prior values. The capture references are passed in the assignment’s shape logic.

The ppx never shares a capture between calls, but a reference made by hand with Shape.get_variable_ref can be passed through ~capture_dims to several operations. It then names one dimension, whichever order shape inference solves them in (gh-ocannl-1545): until inference closes, their dimension variables are equated, as are their row variables, and a dimension with a row has the row’s product. Across a closure (say a Train.forward_once between the operations) the reference carries its solved value: a dimension’s size, and a row variable’s axes (gh-ocannl-1689) together with their product. A row captured again after the closure is the same row – after 2, 3, an operand of 3, 2 or 6 is refused – and a dimension captured again has the row’s product. A reference whose last binding before the closure was a dimension carries only the number, so a row captured from it later is tied by product, as within one inference. The number keeps its basis (gh-ocannl-2033): a dimension of an rgb axis is captured again as rgb, and a free one, closed to the broadcast top 1_(bcast_if_1), as that top. That tie is not transitive (gh-ocannl-1969): row, dimension, row leaves two rows of one product, not one row, and identifying them takes an explicit set_equal of two row references.

Shape.set_equal on two row references equates the rows themselves, not their products, whatever either reference has been solved to so far (gh-ocannl-1617): rows 2, 3 and 6 do not unify. Product equality is what relates a dimension to a row (set_equal on a dimension and a row reference, or Shape.set_dim on a row reference). A row reference bound before a closure takes part in a Row-Row pair with the row it carried, and two such references are compared at the call: different rows raise there, naming both. A row holding a symbolic extent has no numeric product: after its closure it can be captured again as a row, or set_equal-ed with a row reference, but a capture as a dimension, set_dim or a dimension set_equal is refused, naming the closure (gh-ocannl-1655). And a reference made by Shape.get_variable_ref that set_equal relates to a row reference before it is bound is a number, the row’s product: binding it to a row afterwards is refused, since the rows would be tied by product only; pass it to ~capture_dims first.

Capturing requires the specification to be a literal string: every capture pattern in tensor/ppx_op.ml and tensor/ppx_cd.ml matches on a string constant, so x ++ "ab => a" ["b"] works while let s = "ab => a" in x ++ s ["b"] does not expand as a capture.

Einsum gotchas and idioms

Recurring stumbling points, each verified against the sources:

Further features of the syntax extension %cd

Referencing arrays: tensor value, tensor gradient, merge buffer of a tensor node

The %cd syntax uses record-style notation to point to:

The accessor .value can (almost?) always be dropped: by default, tensors in the %cd syntax refer to their value nodes. The forward and backprop code accesses manage roots (via the Tensor.consume_forward_code and Tensor.consume_backprop_code functions).

For example, in a data-parallel computation, gradients of the same param p can be merged across devices using the code p.grad =+ p.grad.merge, combined with an explicit device-to-device transfer.

Block comments

The %cd syntax uses the prefix operator (~~) in a semicolon sequence to introduce block comments:

type Assignments.t =
  ...
  | Block_comment of string * t
  ...

Schematic example: ~~("space" "separated" "comment" "tensor p label_name:" p; <scope of the comment>). The content of the comment uses application syntax, must be composed of strings, <tensor>, <tensor>.value (equivalent to <tensor>), <tensor>.grad components, where <tensor> is any tensor expression or tensor identifier. A tensor component renders as its value node’s Ir.Tnode.label_name (with .grad appended for a .grad component), so the comment does not depend on which code names a renderer minted first.

This syntax used to be very important, because comments in assignments are used to derive file names for generated code. Now, the %cd syntax automatically introduces block comments for code at let-binding points, using the identifier. Currently the comment does not yet incorporate any tensor node labels – and for that reason we are not yet adding comments around function bodies if a function is annotated with %cd. Moreover, we only automatically add comments for code, not for tensors – so the ~~ syntax is still helpful when the comment needs to be more precise for debugging or naming purposes, or when %cd is not used with a let binding, or when we want to pass a forward code directly instead of let-binding it. If an explicit comment is provided at the let-binding level, the automatic one is omitted.

Further features of the syntax extension %op

Name from binding

When an extension point is applied to a let-binding, e.g. let%op mlp_layer ~label ~hid_dim () x = relu ({ w } * x + { b; o = [ hid_dim ] }), it uses the name of the binding (mlp_layer in the example) for the label of the primary tensor created by the extension, if any. This is why the resulting layer tensor in the example has its label starting with "mlp_layer". If the extension is over a semicolon-separated sequence of expressions, the primary tensor can only be in the last component of the sequence, other syntax constructs are handled analogously.

The example let%op mlp_layer ~label ~hid_dim () x = relu ({ w } * x + { b; o = [ hid_dim ] }) also illustrates providing an additional string list to populate labels: label must be of type string list (see The label parameter for which tensors take it).

Label from function argument

When the function directly under %op has no unit parameter (), the resulting (primary) tensor’s label also incorporates the labels of its arguments: in let%op scaled ~gamma x = gamma *. x, the result’s label includes the labels of the applied gamma and x. An unlabelled parameter is always such a label source. A labelled one is when %op sees evidence that it is a tensor: it is a direct operand of a tensor operator in the body (gamma *. x), or it is annotated (~(gamma : Tensor.t), for a tensor the body only passes on). Any other labelled parameter is left alone and may be anything, e.g. let%op scaled ~scale x = x *. !.scale with a float scale (gh-ocannl-1885); so a labelled tensor that the body only hands to another function contributes its label only when annotated. An optional parameter (?scale, ?(stride = 2)) is not a label source and may be anything (gh-ocannl-1787). If the function has a unit parameter (), like mlp_layer above, no tensor parameter is a label source: the label comes from the binding’s name, followed by the label parameter before (), if present.

The label parameter

A parameter with the label label – ~label, ?label or ?(label = [ ... ]) – is a string list that %op puts into the labels of the tensors the function creates per application (gh-ocannl-1987):

label is not otherwise reserved: a ~label that is a tensor by the evidence of the previous section (an operand use, a Tensor.t annotation) is a tensor, and an optional one with that evidence (?(label = fallback) used as an operand) is left alone like any other optional parameter.

Parameters that need a unit parameter, or a label

Two uses of a parameter cannot work where a new user is likely to write them, and %op refuses them with an error naming the fix (gh-ocannl-1987):

Configuring inline declarations: inline output dimensions, initial values

In the %op syntax, inline declarations use record syntax with additional fields to configure the tensor:

A very simple example from micrograd_demo: Micrograd README basic example:

  let%op c = { a = [ -4 ] } + { b = [ 2 ] } in
  ...

How does it relate to let%op c = { a = -4 } + { b = 2 } in ...? Without brackets, the number is used to initialize all cells of the tensor value, and shape inference decides the shape of the tensor. With brackets, the bracketing specifies both all the cells and the exact shape of the tensor.

Need to lift the applications of configuration arguments (up to the unit parameter)

If you recall, inline declared param tensors get lifted out of functions to be defined at the point of a unit () parameter. Our example let%op mlp_layer ~label ~hid_dim () x = relu ({ w } * x + { b; o = [ hid_dim ] }) translates as:

let mlp_layer ~label ~hid_dim () =
  let w = TDSL.param ~more_label:label "w" () 
  and b = TDSL.param ~more_label:label ~output_dims:[ hid_dim ] "b" () in
  fun x -> TDSL.O.(relu (w * x + b))

For this to work properly, when employing such network blocks, their params also need to be introduced at the right moment. At one point, we tried to do this automatically by the %op syntax, but that was confusing to use. So you need to ensure scoping manually. Consider:

(* FIXME: this is wrong! Doesn't bind the parameters at the right place. *)
let%op three_layer_perceptron ~label ~dim1 ~dim2 ~dim3 () x =
  mlp_layer ~label:[ "L3" ] ~hid_dim:dim3 ()
    (mlp_layer ~label:[ "L2" ] ~hid_dim:dim2 ()
       (mlp_layer ~label:[ "L1" ] ~hid_dim:dim1 () x))

This example would work if we used direct inline definitions, but it does not work when the definitions are indirectly in the functions called. We need to write instead:

let three_layer_perceptron ~label ~dim1 ~dim2 ~dim3 () =
  let layer3 = mlp_layer ~label:[ "L3" ] ~hid_dim:dim3 ()
  and layer2 = mlp_layer ~label:[ "L2" ] ~hid_dim:dim2 ()
  and layer1 = mlp_layer ~label:[ "L1" ] ~hid_dim:dim1 () in
  fun x -> layer3 (layer2 (layer1 x))

The manual approach naturally extends to programmatic network architectures:

let mlp ~label ~hid_dims () =
  let layers =
    List.mapi hid_dims ~f:(fun i hid_dim ->
        mlp_layer ~label:[ "L" ^ Int.to_string i ] ~hid_dim ())
  in
  fun x -> List.fold layers ~init:x ~f:(fun x layer -> layer x)

Automatic shape schemes

A %op function gets a shape scheme automatically (gh-ocannl-1751), on top of the explicit layer Tensor.Scheme (gh-ocannl-404, see shape inference), in checking mode:

To see the scheme a closure uses, build it under Tensor.Scheme.traced_by, which returns the schemes the closures built while it ran traced – the ones their applications are checked against – and print one with Tensor.Scheme.to_string (gh-ocannl-1801). A scheme’s results are what the traced body returns: a tensor it made (or a copy of its record) or one of its arguments, or a tuple of them; any other value leaves the scheme without results, printed _. The checks relate the arguments only. Tensor.Scheme.check_definition returns the schemes of its trial the same way, and nn_blocks_definitions.expected prints those of every schemed Nn_blocks block, an inventory of the library’s block signatures. Tensor.Scheme.trace traces a function explicitly, again. The newcomer tutorial Reusable Tensor Programs prints one, shows both kinds of report, and recovers a REPL session from a failed construction.

Which functions qualify. The returned function – right after (), or the function a unit body ends in, through let, ; and let open, as in Nn_blocks.mlp – must take only tensors, or optional tensors. The ppx has no types, so it decides syntactically:

The annotations are type-checked: by OCaml at the parameter, and again by the expansion, which passes the arguments to the scheme as Tensor.t values. A parameter that is only passed on (fun ~train_step x -> bn ~train_step x) gives the ppx no evidence, so annotate it.

Option-valued tensors: one scheme per presence pattern. A function with option-valued tensor parameters gets one scheme per presence pattern, meaning which of them are Some: at most 2^k for k such parameters. Each is traced, checked and refused on its own, under the label <function> [p, no q]. This is how nn_blocks schemes the train/eval switch. ~(train_step : Tensor.t option) is the training step counter as a tensor (Some !@step_n trains, None evaluates), so dropout, multi_head_attention ... () ~train_step ?mask x, the transformer blocks and the batch norms get a training scheme and an evaluation scheme (four for attention, with and without a mask). The closure built at () traces every pattern, so an error in a variant that does not depend on the arguments – a training-only one too – raises when the closure is built, naming the variant. A function without a () stage traces every pattern at its first application. After a session reset each pattern re-traces at its own next application. The traces cost 2^k body runs per closure; opt out if that is too many.

A function with a parameter that is not a tensor – fun ~train_step x with train_step unannotated, a callback, a configuration value – stays unschemed, and expands exactly as before: its parameter’s value does not exist yet when the closure is built, and the body may branch on it. So do let%op rec and function-case bodies. An unlabelled parameter that is not a tensor, against the convention, is a compile error; opt out.

Nested applications. A checked application vouches for the applications in its body (gh-ocannl-1814): its scheme holds their constraints, so they run without a check of their own, and an argument error in one is reported by the outer check, naming the outer function. A function’s first application still traces its other presence patterns, which that scheme does not hold. Under an opted-out or refused function, nothing vouches for them and they are checked as usual. That is why the nn_blocks stacks (transformer_encoder, transformer_decoder, decoder_only) are schemed like their layers: an application is one check, whose scheme gives the whole model’s signature, for a trace per presence pattern, each a fraction of one build. sokoban_cnn and mobile_cnn are schemed the same way.

Refused traces degrade, with a warning. A function whose trace a scheme cannot carry – Shape.set_dim or the like on a capture reference made outside it, forcing dims or projections, a body that cannot run on tracers – stays unschemed, once the trace’s rewind succeeded. It is a valid function, and its applications report what they always reported. The refusal is printed to stderr once per function and reason in a process (OCANNL warning: the shape scheme of <function> was refused ...), with the reason rendered as a shape error is. Opting out, op_shape_schemes=off, and a closure built inside another function’s trace (built again, and traced, when that function is applied) print nothing.

A capture reference made outside the function – at its () stage, say, and passed to an operation’s ~capture_dims in the returned function – is no refusal (gh-ocannl-1791): it is an input of the scheme, related at each application to what earlier applications bound it to, so an application conflicting with them is refused at that application, before the body runs.

Random draws in the body are no refusal either (gh-ocannl-1890), default-initialized { w } declarations included: the scheme checks the shapes around a draw, and the body’s own registration sizes it. Only the explicit Tensor.Scheme.instantiate, which has no body to run, refuses a scheme with draws.

Side effects. The trace runs the body once more per closure (once per presence pattern with option-valued tensors). Under the rewind it leaves no shape state, no tensor and no session id behind (traces and checks allocate tensor ids and shape-variable ids from reserved ranges, so printed ids and shapes, compilation choices, generated code and random draws are the same with or without schemes; node uids come from the one session counter, so the trace’s discarded nodes shift later uids, which nothing but node identity and a few uid-spelled identifiers – MMA fragment registers, recompute_<uid> routine names – reads, gh-ocannl-2053), but OCaml side effects of the body (a counter, a print) happen.

Opting out. [@@ocannl.no_scheme] on the binding, or [@ocannl.no_scheme] on a fun expression, leaves a function unschemed: for a body whose constraints depend on something other than its closure and its arguments’ shapes (mutable state, say), or whose side effects must run exactly once per application. The configuration key op_shape_schemes=off switches the schemes off for the whole process.

Checking a definition early

A block’s () trace reports an argument-independent error only when the closure is built. To catch it at the definition, apply the definition to a trial or default configuration with Tensor.Scheme.check_definition, which runs the trial in a window that is always rewound and allocates no session id, and names the definition in its error:

```ocaml doc-check=op_schemes_check_definition open Ocannl.Nn_blocks.DSL_modules

let%op block ~hid_dim () x = relu (({ w } * x) + { b = 0.; o = [ hid_dim ] }) let () = ignore (Tensor.Scheme.check_definition ~name:“block” (fun () -> ignore (block ~hid_dim:8 ())))


It catches the errors that hold for every configuration, and the configuration-dependent ones of the trial configuration only; it costs one trace (and the trial's construction) at program start. Prefer it to a bare `let _ = block conf ()`, which leaves the trial's parameters registered in the session. It traces every presence pattern of the closures the trial builds at `()`, so a training-only error is caught too. It does not reach the functions that stay unschemed, beyond the sub-blocks their trial builds, nor a function without a `()` stage that the trial only builds (that one traces at its first application). `lib/nn_blocks.ml` does not run it at load time; `test/operations/nn_blocks_definitions` runs it over every schemed block.

## The syntax extension %%extend_dsls

This syntax extension creates a module `DSL_modules` with the same submodules as `Operation.DSL_modules`. It removes the boilerplate associated with introducing new operators into the modules `TDSL`, `NTDSL` and their `O` submodules. The payload (i.e. content) of `%%extend_dsls` must be non-recursive let-bindings. They are parsed using a slight variant of the `%op` syntax, and are inserted into the DSL modules. The identifiers of the root operator functions of the definitions, if unqualified, are prefixed with the appropriate module, similarly to the behavior of inline definitions. Initialization expressions remain forward-only.

The DSL modules expose the value `grad_spec` that can be useful for defining operators via a "scheme" function. See the example using the `box_muller` helper at the beginning of `lib/nn_blocks.ml`. The definitions there use the `%oc` escape extension to avoid the prefixing mentioned above.

## Implementation details

### The hard-coded to-the-power-of operator

OCANNL has a built-in numerical binary operation to-power-of: `Ops.ToPowOf`. As part of assignments, the corresponding operator is `**`. Here is the full definition of the to-power-of tensor operation from [Operation](../tensor/operation.ml):

```ocaml
let rec pointpow ?(label : string list = []) ~grad_spec p t1 : Tensor.t =
  let module NTDSL = struct
    include Initial_NTDSL

    module O = struct
      include NDO_without_pow

      let ( **. ) ?label base exp = pointpow ?label ~grad_spec:Tensor.Prohibit_grad exp base
    end
  end in
  let p_t = NTDSL.number p in
  let%cd op_asn ~t ~t1 ~t2 ~projections = v =: v1 ** v2 ~projections in
  let%cd grad_asn =
    if Tensor.is_prohibit_grad grad_spec then fun ~v:_ ~g:_ ~t1:_ ~t2:_ ~projections:_ -> Asgns.Noop
    else if Float.equal p 2.0 then fun ~v:_ ~g ~t1 ~t2:_ ~projections -> g1 =+ p_t *. t1 * g
    else if Float.equal p 1.0 then fun ~v:_ ~g ~t1 ~t2:_ ~projections -> g1 =+ g
    else fun ~v:_ ~g ~t1 ~t2:_ ~projections -> g1 =+ p_t *. (t1 **. (p -. 1.)) * g
  in
  Tensor.binop ~label:("**." :: label) ~compose_op:Pointwise_bin ~op_asn ~grad_asn ~grad_spec t1 p_t

On the Tensor level, this is implemented as a binary tensor operation, but it is exposed as a unary tensor operation! To avoid the complexities of propagating gradient into the exponent, Operation.pointpow is implemented as a function of only one tensor, the exponent is a number. We hard-code the pointwise-power-of operator NTDSL.O.( **. ), resp. TDSL.O.( **. ), in the %cd and %op syntaxes, to pass the numeric value to pointpow (the second argument of **.) without converting it to a tensor first.

Intricacies of the syntax extension %cd

The syntax %cd translator needs to accomplish more than a context-free conversion of a concise notation to an Assignments.comp data-type. In particular:

Embedded nodes

In fact, the syntax %cd produces Assignments.comp values:

type comp = {
  asgns : t;
  embedded_nodes : Set.M(Tnode).t;
}

The tensor nodes that are in asgns but not in embedded_nodes, and are on-device, must already be present in contexts with which the computation is linked. Such non-embedded nodes can be seen as inputs to the computation – except that for backprop code of a tensor, they are actually the outputs! Embedded nodes are closely related to rootness – when a node has not been used in the code of another tensor, it is a root (a forward root for value nodes and a backprop root for grad nodes). embedded_nodes were roots the first time they were used in asgns. Parameters, as created by Tensor.param, are not embedded in the code that uses them and thus will not be in embedded_nodes of the forward and backprop code over the parameters; however, they will constitute the embedded_nodes of the Tensor.init_params code.