Module Ocannl_tensor.Operation

Helper to compute compose_op: if spec is provided, use Einsum; otherwise use Pointwise_bin.

Helper to compute transpose_op: if spec is provided, use Permute; otherwise use Pointwise_un.

An identity operation whose sole purpose is to give the result a tensor node distinct from the argument's: the assignment-level precision conversion then happens between the two nodes when their precisions differ. This is the "cast twin" primitive of the mixed-precision recipes (gh-ocannl-492): pin the result at a reduced precision to make it a low-precision copy of an f32 master weight; the gradient accumulates back through the cast, widening (e.g. f16 -> f32) at the accumulating assignment.

An identity operation whose result's rows resolve at the use site (gh-ocannl-544): the result widens to what its consumers demand, and the argument broadcasts into it. Operation results otherwise close down to their arguments' shapes, so this is the explicit request for use-site widening — e.g. stretch !.1.0 is a shape-inferred constant 1 usable as an einsum operand of any shape (formerly spelled 0.5 + 0.5).

Ternary select: computes where(pred, a, b) element-wise. When ~spec is given it routes through ternary einsum: over the product space of all spec labels, at each point the predicate element selects the branch-2 or branch-3 element (select happens *before* accumulation). Labels absent from the LHS become reduction axes.

Gradient: flows to a masked by the predicate and to b masked by its complement; no gradient to the predicate.

Documented caveat: when a reduced axis is absent from a branch (branch broadcasts along it), the broadcast value is accumulated once per selected reduction iteration ("count-weighted") — this is the honest semantics of reduce-after-select. For example, with spec "p;a;b=>out" and a broadcasting along a reduced axis, a's scalar value is summed once for each reduction step where the predicate is 1.

Similar to the explicit mode of numpy.einsum, the binary variant. Can compute various forms of matrix multiplication, inner and outer products, etc.

Note that "a,b->c" from numpy is "a;b=>c" in OCANNL, since "->" is used to separate the input and the output axes.

Like einsum, but adds instead than multiplying the resulting values.

Ternary einsum: contracts three tensors under a three-RHS spec rhs1 ; rhs2 ; rhs3 => lhs. The forward pass computes v[lhs] =+ mul3(v1[rhs1], v2[rhs2], v3[rhs3]) over the product space of all axis labels. Axis labels present in any RHS but absent from the LHS become reduction axes.

WARNING: for chain-structured contractions such as "ij;jk;km=>im" the fused kernel iterates the full product space (O(N⁴) for N-dimensional operands), which is less efficient than a binary chain (O(N³)). Use ternary einsum for patterns where all three operands share reduction axes, e.g. bilinear forms.

Similar to the explicit mode of numpy.einsum, the unary variant. Can permute axes, extract diagonals, compute traces etc.

Note that "a->c" from numpy is "a=>c" in OCANNL, since "->" is used to separate the input and the output axes.

The axis kind along which stack introduces the fresh leading axis.

Stack rhses along a fresh leading axis of kind axis (result rank = operand rank + 1, true stacking rather than concatenation along an existing axis). Each operand is unsqueezed with a size-1 axis via einsum1, then the unsqueezed intermediates are concatenated along that new axis via concat. The expansion and concatenation einsum specs are generated on the fly from the operand count.

Max-reduction einsum. The gradient gate lives in the product space of the operation — one bit per (result position, reduced position) pair — so it is exact even when reduction windows overlap (e.g. convolution-style specs with stride < window). Ties gate the full upstream gradient to every achieving position (duplication, not a normalized split).

nonoverlapping (gh-ocannl-527): the caller asserts that each RHS1 position feeds at most one result position — per convolution axis, stride at least the window extent (a full reduction trivially qualifies). On that domain the gate contracts to one bit per input position (cond_rhs1, the pre-gh-512 formulation): every input position belongs to at most one window, so no gate-bit collisions exist and the two formulations agree exactly, ties included — this is a cost specialization, not a semantics choice. The product-space intermediates it avoids cost 1.8-2.6x on the non-overlapping conv-pooling benchmarks. Outside the asserted domain the cheap gate is last-write-wins across colliding windows, i.e. wrong — the flag is a contract, not a heuristic; when in doubt keep the default.

This generalizes the tropical matrix multiplication to arbitrary indices combinations.

The gradient gate lives in the product space of the operation — one bit per (result position, contracted position) pair — so both gradients are exact for arbitrary specs, including overlapping reduction windows (stride < window) and RHS2 indices independent of RHS1. Ties gate the full upstream gradient to every achieving pair (duplication, not a normalized split).

nonoverlapping (gh-ocannl-527): the caller asserts the einmax1 domain condition — each RHS1 position feeds at most one result position — and additionally that the (result, contracted) pair is recoverable from the RHS1 position (RHS1's index space covers RHS2's, the convolution-like shape: kernel indices contracted against strided input indices; the pre-gh-512 g2 limitation delimits exactly this). On that domain the gate contracts to one bit per input position and both gradients agree with the product-space gate exactly, ties included — a cost specialization, not a semantics choice (1.8-2.6x on the non-overlapping conv-pooling benchmarks). Outside it the cheap gate is wrong (last-write-wins collisions, misattributed g2); when in doubt keep the default.

A fully-shape-inferred tensor that is initialized with the offset of each cell.

range is a 1D tensor of shape upto, spans 0 inclusive, upto exclusive.

A stop_gradient is an identity in the forward pass and a no-op in the backprop pass.

Generates a single uniform random number using a counter symbol for PRNG state. This is useful for sequential sampling in recurrent contexts.

Note: unlike uniform, this stays pointwise (no pin_counter_spec): the user counter's shape must propagate to the result through the key operand, so the key slot cannot be reduced away. Entropy across draws comes from the user counter; within one draw, range_over_offsets may close smaller than the result and repeat values along broadcast axes.

DEPRECATED (gh-ocannl-509): a wasteful variant of uniform that produces a single value from each 4x32 random bits. The packed uniform is now total over shapes (round-up counter inference, tail-peeled lowering) and virtualizes via lane extraction, so this shape-safety fallback is no longer needed; it remains as the IR-internal substrate and for reproducing pre-0.9 random streams. The bit-spreading in int32_to_uint4x32/uint32_to_uint4x32 ensures good entropy even with the 2-round "light" threefry variant.

A centered uniform distribution over [-scale/2, scale/2)] built from the packed {!uniform}. With the default [scale = 0.5] this is the default parameter initialization (see {!Make_DSL.default_param_init}).

DEPRECATED (gh-ocannl-509): use centered_uniform_param_init. Kept for reproducing pre-0.9 random streams.

DEPRECATED (gh-ocannl-509): use default_uniform_param_init.

A wasteful variant of uniform_at that produces a single value from each 4x32 random bits. The bit-spreading in int32_to_uint4x32/uint32_to_uint4x32 ensures good entropy even with the 2-round "light" threefry variant.

The input i dimensions default to empty. The batch and output dimensions will be inferred if omitted. Note: the data should have no padding and if padding is inferred, the data will be copied; otherwise, the resulting tensor value shares host memory with the ndarray.

The dimensions are taken from the provided ndarray, but the split into axis kinds still needs to be inferred (or provided). Assumes no padding. Input axes are not inferred (empty if omitted). See also: reshape and TDSL.wrap_param.

Assumes the ndarray is padded as given. This means the dimensions of the ndarray will differ from the dimensions of the tensor by the padding. See also: TDSL.wrap.

The output dimensions are taken from the provided ndarray, assuming precisely the first axis is a batch axis, assumes no input axes and the batch dimensions are inferred. Empty output dimensions are allowed and represent scalars. Assumes the data has no padding, and data is copied if padding is inferred. See also: reshape and wrap.

Creates a tensor by initializing values using a function from indices to values. The dimensions are split into axis kinds as specified, there is no shape inference. Recall that input axes are rightmost.