Skip to content

Graph & operators

The lazy Tensor graph and the operators it is built from. These are exposed at the top level, so aneforge.graph.conv is reached as af.conv.

graph

Lazy Tensor graph, op constructors, and nn helpers (conv, attention, GEGLU). The device-free frontend; _compile.py lowers it.

Tensor

A node in the compute graph.

Source code in aneforge/graph.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
class Tensor:
  """A node in the compute graph."""

  __slots__ = ("shape", "op", "srcs", "attrs", "_name")

  def __init__(self, shape: Sequence[int], op: str, srcs: Sequence["Tensor"] = (),
        attrs: dict | None = None) -> None:
    self.shape = tuple(int(d) for d in shape)
    # ANE dimension model is rank <= 5; rank-6+ fails ANECCompile. Guard early.
    if len(self.shape) > 5:
      raise ValueError(
        f"aneforge: tensor rank {len(self.shape)} exceeds the ANE maximum of 5 "
        f"(op={op!r}, shape={self.shape}); the ANE dimension model is rank<=5.")
    self.op = op
    self.srcs = list(srcs)
    self.attrs = attrs or {}
    self._name = ""

  # -- elementwise ------------------------------------------------------- #
  # param-free unary ops: graph op name == MIL op name (see _compile._e_unary)
  def relu(self) -> "Tensor": return Tensor(self.shape, "relu", [self])
  def gelu(self) -> "Tensor": return Tensor(self.shape, "gelu", [self])
  def silu(self) -> "Tensor": return Tensor(self.shape, "silu", [self])
  def sigmoid(self) -> "Tensor": return Tensor(self.shape, "sigmoid", [self])
  def tanh(self) -> "Tensor": return Tensor(self.shape, "tanh", [self])
  def exp(self) -> "Tensor": return Tensor(self.shape, "exp", [self])
  def sqrt(self) -> "Tensor": return Tensor(self.shape, "sqrt", [self])
  def abs(self) -> "Tensor": return Tensor(self.shape, "abs", [self])
  def square(self) -> "Tensor": return Tensor(self.shape, "square", [self])
  def sin(self) -> "Tensor": return Tensor(self.shape, "sin", [self])
  def cos(self) -> "Tensor": return Tensor(self.shape, "cos", [self])
  def erf(self) -> "Tensor": return Tensor(self.shape, "erf", [self])
  def softplus(self) -> "Tensor": return Tensor(self.shape, "softplus", [self])
  def relu6(self) -> "Tensor": return Tensor(self.shape, "relu6", [self])
  def softsign(self) -> "Tensor": return Tensor(self.shape, "softsign", [self])
  def atan(self) -> "Tensor": return Tensor(self.shape, "atan", [self])
  def exp2(self) -> "Tensor": return Tensor(self.shape, "exp2", [self])

  # unary ops carrying a parameter
  def log(self, eps: float = 0.0) -> "Tensor": return Tensor(self.shape, "log", [self], {"eps": eps})
  def rsqrt(self, eps: float = 0.0) -> "Tensor": return Tensor(self.shape, "rsqrt", [self], {"eps": eps})
  def inverse(self, eps: float = 0.0) -> "Tensor":
    """`1 / x` (elementwise reciprocal); `eps` floors the divide."""
    return Tensor(self.shape, "inverse", [self], {"eps": eps})
  def elu(self, alpha: float = 1.0) -> "Tensor": return Tensor(self.shape, "elu", [self], {"alpha": alpha})
  def leaky_relu(self, alpha: float = 0.01) -> "Tensor": return Tensor(self.shape, "leaky_relu", [self], {"alpha": alpha})
  def clip(self, lo: float, hi: float) -> "Tensor": return Tensor(self.shape, "clip", [self], {"lo": lo, "hi": hi})

  # parametric activations (scalar fp16 params, like clip/elu)
  def scaled_tanh(self, alpha: float = 1.0, beta: float = 1.0) -> "Tensor":
    """`alpha * tanh(beta * x)`."""
    return Tensor(self.shape, "scaled_tanh", [self], {"alpha": alpha, "beta": beta})
  def threshold(self, alpha: float = 0.0) -> "Tensor":
    """`max(x, alpha)`."""
    return Tensor(self.shape, "threshold", [self], {"alpha": alpha})
  def thresholded_relu(self, alpha: float = 1.0) -> "Tensor":
    """`x if x >= alpha else 0`."""
    return Tensor(self.shape, "thresholded_relu", [self], {"alpha": alpha})
  def clamped_relu(self, alpha: float = 0.0, beta: float = 6.0) -> "Tensor":
    """`min(beta, x)` for `x >= 0` else `min(beta, alpha * x)` (a leaky relu6)."""
    return Tensor(self.shape, "clamped_relu", [self], {"alpha": alpha, "beta": beta})
  def sigmoid_hard(self, alpha: float = 0.2, beta: float = 0.5) -> "Tensor":
    """Hard sigmoid: `min(max(alpha * x + beta, 0), 1)`."""
    return Tensor(self.shape, "sigmoid_hard", [self], {"alpha": alpha, "beta": beta})
  def linear_activation(self, alpha: float = 1.0, beta: float = 0.0) -> "Tensor":
    """`alpha * x + beta` (scalar affine, fused as one op)."""
    return Tensor(self.shape, "linear_activation", [self], {"alpha": alpha, "beta": beta})

  def greater(self, o) -> "Tensor":
    """Elementwise `x > o` -> a BOOL tensor (use as the `cond` of `af.select`)."""
    if not isinstance(o, Tensor):
      raise TypeError("greater expects a graph Tensor (compare against a streamed weight via select)")
    return Tensor(_broadcast(self.shape, o.shape), "greater", [self, o])

  def _cmp(self, o, op: str) -> "Tensor":
    if not isinstance(o, Tensor):
      raise TypeError(f"{op} expects a graph Tensor (compare against a streamed weight via select)")
    return Tensor(_broadcast(self.shape, o.shape), op, [self, o])

  def less(self, o) -> "Tensor": return self._cmp(o, "less")               # x < o  -> bool
  def equal(self, o) -> "Tensor": return self._cmp(o, "equal")             # x == o -> bool
  def not_equal(self, o) -> "Tensor": return self._cmp(o, "not_equal")     # x != o -> bool
  def less_equal(self, o) -> "Tensor": return self._cmp(o, "less_equal")   # x <= o -> bool
  def greater_equal(self, o) -> "Tensor": return self._cmp(o, "greater_equal")  # x >= o -> bool
  def logical_not(self) -> "Tensor": return Tensor(self.shape, "logical_not", [self])  # ~bool

  def floor(self) -> "Tensor": return Tensor(self.shape, "floor", [self])
  def ceil(self) -> "Tensor": return Tensor(self.shape, "ceil", [self])
  def round(self) -> "Tensor": return Tensor(self.shape, "round", [self])
  def sign(self) -> "Tensor": return Tensor(self.shape, "sign", [self])

  def prelu(self, alpha) -> "Tensor":
    """Per-channel PReLU: `x if x>0 else alpha[c]*x`. `alpha`: [C]; input rank>=3 [N,C,...]."""
    alpha = np.asarray(alpha)
    if len(self.shape) < 3:
      raise ValueError(f"prelu needs rank>=3 [N,C,...]; got {self.shape}")
    if alpha.shape != (self.shape[1],):
      raise ValueError(f"prelu alpha {alpha.shape} != channels {self.shape[1]}")
    return Tensor(self.shape, "prelu", [self], {"alpha": alpha})

  def __pow__(self, o) -> "Tensor": return _binary(self, o, "pow")
  def pow(self, o) -> "Tensor": return _binary(self, o, "pow")

  def reverse(self, axes) -> "Tensor":
    """Reverse along `axes` (native `reverse`)."""
    if not isinstance(axes, (tuple, list)):
      axes = (axes,)
    axes = tuple(a % len(self.shape) for a in axes)
    return Tensor(self.shape, "reverse", [self], {"axes": axes})

  def tile(self, reps) -> "Tensor":
    """Repeat `reps[i]` times along each axis (native `tile`; factors of {2,3,4,8})."""
    reps = tuple(int(r) for r in reps)
    out = tuple(d * r for d, r in zip(self.shape, reps))
    return Tensor(out, "tile", [self], {"reps": reps})

  def reduce_log_sum_exp(self, axes) -> "Tensor":
    """log(sum(exp(x))) over `axes` (native, stable softmax denominator)."""
    return self._reduce("reduce_log_sum_exp", axes)

  def clamp(self, lo: float, hi: float) -> "Tensor": return self.clip(lo, hi)  # alias

  def __add__(self, o) -> "Tensor": return _binary(self, o, "add")
  def __sub__(self, o) -> "Tensor": return _binary(self, o, "sub")
  def __truediv__(self, o) -> "Tensor": return _binary(self, o, "real_div")

  def __mul__(self, o) -> "Tensor":
    if isinstance(o, (int, float)):
      return Tensor(self.shape, "muls", [self], {"k": float(o)})
    return _binary(self, o, "mul")
  __rmul__ = __mul__

  def adds(self, k: float) -> "Tensor":
    """`x + scalar` as a fused scalar-add (the only fused way to inject a scalar offset; `+` needs two Tensors)."""
    return Tensor(self.shape, "adds", [self], {"k": float(k)})

  # -- linear algebra ---------------------------------------------------- #
  def __matmul__(self, W) -> "Tensor":
    """`x @ W`. `W` is a streamed weight array, or a Tensor for an activationxactivation product."""
    if isinstance(W, Tensor):
      if self.shape[-1] != W.shape[-2]:
        raise ValueError(f"matmul: {self.shape} @ {W.shape} shape mismatch")
      return Tensor(self.shape[:-1] + (W.shape[-1],), "bmm", [self, W])
    W = np.asarray(W); _check_dtype(W, "matmul weight")
    if W.ndim != 2 or self.shape[-1] != W.shape[0]:
      raise ValueError(f"matmul: x{self.shape} @ W{W.shape} shape mismatch")
    # store as [N,K] and consume with transpose_y=true (the proven int8 layout)
    return Tensor(self.shape[:-1] + (W.shape[1],), "matmul", [self],
           {"wt": np.ascontiguousarray(W.T)})

  def bmm_weight(self, W) -> "Tensor":
    """Batched matmul `self [B,M,K] @ W [B,K,N]`, `W` a baked QUANTIZABLE weight (unlike `@` with a `_const`,
    which stays fp16) -- for per-expert projections, so `compress=` shrinks them."""
    W = np.asarray(W); _check_dtype(W, "bmm_weight")
    if W.ndim != 3 or self.shape[-1] != W.shape[-2]:
      raise ValueError(f"bmm_weight: x{self.shape} @ W{W.shape} shape mismatch")
    return Tensor(self.shape[:-1] + (W.shape[-1],), "bmm_w", [self], {"wt": np.ascontiguousarray(W)})

  def linear(self, W, bias=None) -> "Tensor":
    """`x @ W.T (+ bias)`. `W` is [out, in] (PyTorch convention)."""
    W = np.asarray(W); _check_dtype(W, "linear weight")
    if W.ndim != 2 or self.shape[-1] != W.shape[1]:
      raise ValueError(f"linear: x{self.shape} W{W.shape} shape mismatch")
    attrs: dict[str, Any] = {"wt": np.ascontiguousarray(W)}
    if bias is not None:
      attrs["bias"] = np.asarray(bias).astype(np.float32)
    return Tensor(self.shape[:-1] + (W.shape[0],), "matmul", [self], attrs)

  def transpose(self, perm) -> "Tensor":
    perm = tuple(p % len(self.shape) for p in perm)
    return Tensor(tuple(self.shape[p] for p in perm), "transpose", [self], {"perm": perm})

  def reshape(self, *shape) -> "Tensor":
    if len(shape) == 1 and isinstance(shape[0], (tuple, list)):
      shape = tuple(shape[0])
    return Tensor(tuple(shape), "reshape", [self])

  def squeeze(self, axes) -> "Tensor":
    """Remove size-1 dims at `axes` (native `squeeze`)."""
    if not isinstance(axes, (tuple, list)):
      axes = (axes,)
    axes = tuple(a % len(self.shape) for a in axes)
    for a in axes:
      if self.shape[a] != 1:
        raise ValueError(f"squeeze: axis {a} has size {self.shape[a]} != 1")
    out = tuple(d for i, d in enumerate(self.shape) if i not in axes)
    return Tensor(out, "squeeze", [self], {"axes": axes})

  def expand_dims(self, axes) -> "Tensor":
    """Insert size-1 dims at `axes` (native `expand_dims`); axes index the OUTPUT rank."""
    if not isinstance(axes, (tuple, list)):
      axes = (axes,)
    out_rank = len(self.shape) + len(axes)
    axes = sorted(a % out_rank for a in axes)
    out, src = [], iter(self.shape)
    for i in range(out_rank):
      out.append(1 if i in axes else next(src))
    return Tensor(tuple(out), "expand_dims", [self], {"axes": tuple(axes)})

  def flatten2d(self, axis: int = 1) -> "Tensor":
    """Collapse to 2-D about `axis`: `[:axis]` -> rows, `[axis:]` -> cols (native `flatten2d`)."""
    ax = axis % len(self.shape)
    rows = int(np.prod(self.shape[:ax])) if ax > 0 else 1
    cols = int(np.prod(self.shape[ax:]))
    return Tensor((rows, cols), "flatten2d", [self], {"axis": ax})

  def slice_by_size(self, begin, size) -> "Tensor":
    """Static per-axis slice `x[begin[i]:begin[i]+size[i]]` (native `slice_by_size`)."""
    begin = [int(b) for b in begin]; size = [int(s) for s in size]
    if len(begin) != len(self.shape) or len(size) != len(self.shape):
      raise ValueError(f"slice_by_size: begin/size must have rank {len(self.shape)}")
    for i, (b, s) in enumerate(zip(begin, size)):
      if b < 0 or s <= 0 or b + s > self.shape[i]:
        raise ValueError(f"slice_by_size: axis {i} window [{b}:{b+s}] out of range for {self.shape[i]}")
    return Tensor(tuple(size), "slice_by_size", [self], {"begin": begin, "size": size})

  # -- reductions / normalisation --------------------------------------- #
  def _reduce(self, op: str, axes) -> "Tensor":
    if not isinstance(axes, (tuple, list)):
      axes = (axes,)
    axes = tuple(a % len(self.shape) for a in axes)
    out = tuple(1 if i in axes else d for i, d in enumerate(self.shape))
    return Tensor(out, op, [self], {"axes": axes})

  def mean(self, axes) -> "Tensor": return self._reduce("reduce_mean", axes)
  def sum(self, axes) -> "Tensor": return self._reduce("reduce_sum", axes)
  def amax(self, axes) -> "Tensor": return self._reduce("reduce_max", axes)
  def amin(self, axes) -> "Tensor": return self._reduce("reduce_min", axes)

  def cumsum(self, axis: int = -1) -> "Tensor":
    """Cumulative sum along the last axis as `x @ triu_ones` (no native cumsum)."""
    ax = axis % len(self.shape)
    if ax != len(self.shape) - 1:
      raise ValueError(f"cumsum: only the last axis is supported (got axis={axis}, "
              f"rank {len(self.shape)}); transpose the axis to last first")
    N = self.shape[-1]
    W = np.tril(np.ones((N, N), dtype=np.float32)).astype(np.float16)  # linear(W)=x@W.T=x@triu
    return self.linear(W, None)
  def l1_norm(self, axes) -> "Tensor":
    """`sum(|x|, axes)` (keepdims) via the native `reduce_l1_norm` op."""
    return self._reduce("reduce_l1_norm", axes)
  def log_sum(self, axes) -> "Tensor":
    """`log(sum(x, axes))` (keepdims). Expects a positive input."""
    return self._reduce("reduce_log_sum", axes)
  def sum_square(self, axes) -> "Tensor":
    """`sum(x**2, axes)` (keepdims) via the native `reduce_sum_square` op."""
    return self._reduce("reduce_sum_square", axes)

  def softmax(self, axis: int = -1) -> "Tensor":
    return Tensor(self.shape, "softmax", [self], {"axis": axis % len(self.shape)})

  def l2_norm(self, axis: int = -1, eps: float = 1e-12) -> "Tensor":
    """Per-axis L2-normalize `x / sqrt(sum(x**2, axis) + eps)` (built per-axis; MIL `l2_norm` is all-dims)."""
    ax = axis % len(self.shape)
    return Tensor(self.shape, "l2_norm", [self], {"axis": ax, "eps": float(eps)})

  def argmax(self, axis: int = -1) -> "Tensor":
    """Argmax along `axis` (keepdims). GlobalArgMinMax bridge (cut); 2D [C,W] only, indices fp16-encoded."""
    if len(self.shape) != 2:
      raise ValueError(f"argmax: only 2D [C,W] inputs are supported; got {self.shape}")
    ax = axis % 2
    out = tuple(1 if i == ax else d for i, d in enumerate(self.shape))
    return Tensor(out, "argmax", [self], {"axis": ax})

  def rms_norm(self, gamma, eps: float = 1e-5) -> "Tensor":
    """RMSNorm over the last dim. `gamma`: a [D] array (baked) or a [1,D] Tensor (trainable)."""
    if isinstance(gamma, Tensor):
      xn = self.rms_norm(np.ones(self.shape[-1], np.float32), eps)
      return xn * gamma
    gamma = np.asarray(gamma); _check_dtype(gamma, "rms_norm gamma")
    if len(self.shape) != 2 or gamma.shape != (self.shape[-1],):
      raise ValueError(f"rms_norm expects 2D [M,D] with gamma [D]; got {self.shape}, {gamma.shape}")
    return Tensor(self.shape, "rms_norm", [self], {"gamma": gamma, "eps": float(eps)})

  def layer_norm(self, gamma, beta, eps: float = 1e-5) -> "Tensor":
    """LayerNorm over the last dim (2D [M,D]). `gamma`/`beta`: [D] arrays (baked) or [1,D] Tensors (trainable; pass both)."""
    if isinstance(gamma, Tensor) or isinstance(beta, Tensor):
      if not (isinstance(gamma, Tensor) and isinstance(beta, Tensor)):
        raise TypeError("layer_norm: a trainable affine needs both gamma and beta as Tensors")
      D = self.shape[-1]
      xn = self.layer_norm(np.ones(D, np.float32), np.zeros(D, np.float32), eps)
      return xn * gamma + beta
    gamma = np.asarray(gamma); beta = np.asarray(beta)
    _check_dtype(gamma, "layer_norm gamma"); _check_dtype(beta, "layer_norm beta")
    if len(self.shape) != 2 or gamma.shape != (self.shape[-1],):
      raise ValueError(f"layer_norm expects 2D [M,D] with gamma/beta [D]; got {self.shape}, {gamma.shape}")
    return Tensor(self.shape, "layer_norm", [self], {"gamma": gamma, "beta": beta, "eps": float(eps)})

  def channel_layer_norm(self, gamma, beta, eps: float = 1e-5) -> "Tensor":
    """LayerNorm over the CHANNEL axis of [N,C,1,S] (ANE transformer layout; no transpose). `gamma`/`beta`: [C]."""
    gamma = np.asarray(gamma); beta = np.asarray(beta)
    _check_dtype(gamma, "channel_layer_norm gamma"); _check_dtype(beta, "channel_layer_norm beta")
    if (len(self.shape) != 4 or self.shape[2] != 1
        or gamma.shape != (self.shape[1],) or beta.shape != (self.shape[1],)):
      raise ValueError(f"channel_layer_norm expects [N,C,1,S] with gamma/beta [C]; "
              f"got {self.shape}, {gamma.shape}")
    return Tensor(self.shape, "channel_layer_norm", [self], {"gamma": gamma, "beta": beta, "eps": float(eps)})

  def group_norm(self, gamma, beta, num_groups: int, eps: float = 1e-5) -> "Tensor":
    """GroupNorm over [1,C,H,W]. `gamma`/`beta`: [C] arrays (baked) or [1,C,1,1] Tensors (trainable; pass both)."""
    if isinstance(gamma, Tensor) or isinstance(beta, Tensor):
      if not (isinstance(gamma, Tensor) and isinstance(beta, Tensor)):
        raise TypeError("group_norm: a trainable affine needs both gamma and beta as Tensors")
      C = self.shape[1]
      xn = self.group_norm(np.ones(C, np.float32), np.zeros(C, np.float32), num_groups, eps)
      return xn * gamma + beta
    gamma = np.asarray(gamma); beta = np.asarray(beta)
    _check_dtype(gamma, "group_norm gamma"); _check_dtype(beta, "group_norm beta")
    if len(self.shape) != 4 or self.shape[0] != 1 or self.shape[1] % num_groups:
      raise ValueError(f"group_norm expects [1,C,H,W] with C%groups==0; got {self.shape}, G={num_groups}")
    # tiled lowering reduces over max(C/groups, H*W) vs the 65536 per-axis cap (finding_sd15)
    _, C, H, W = self.shape
    axis = max(C // num_groups, H * W)
    if axis > 65536:
      raise ValueError(
        f"group_norm: largest tiled axis max(C/groups, H*W) = {axis} exceeds "
        f"the ANE per-axis bound (65536) for {self.shape} with groups={num_groups}; "
        f"reduce the feature map or channels, or use more groups.")
    return Tensor(self.shape, "group_norm", [self],
           {"gamma": gamma, "beta": beta, "groups": num_groups, "eps": float(eps)})

  # -- spatial ----------------------------------------------------------- #
  def _pool(self, op: str, k: int, stride: int | None, pad: int,
            ceil_mode: bool = False, exclude_pad: bool = False) -> "Tensor":
    stride = stride or k
    N, C, H, W = self.shape
    r = stride - 1 if ceil_mode else 0                       # ceil((s+2p-k)/stride) vs floor
    out = (N, C, (H + 2 * pad - k + r) // stride + 1, (W + 2 * pad - k + r) // stride + 1)
    return Tensor(out, op, [self], {"k": k, "stride": stride, "pad": pad,
                                    "ceil_mode": ceil_mode, "exclude_pad": exclude_pad})

  def max_pool(self, k: int, stride: int | None = None, pad: int = 0, ceil_mode: bool = False) -> "Tensor":
    return self._pool("max_pool", k, stride, pad, ceil_mode)

  def avg_pool(self, k: int, stride: int | None = None, pad: int = 0,
               ceil_mode: bool = False, exclude_pad: bool = False) -> "Tensor":
    """Average pool. `exclude_pad` divides by the valid (non-pad) cell count (ONNX `count_include_pad=0`); default includes pad cells (divide by full kernel area)."""
    return self._pool("avg_pool", k, stride, pad, ceil_mode, exclude_pad)

  def upsample(self, scale: int = 2) -> "Tensor":
    """Nearest-neighbour upsample [N,C,H,W] -> [N,C,scale*H,scale*W]."""
    N, C, H, W = self.shape
    return Tensor((N, C, H * scale, W * scale), "upsample", [self], {"scale": scale})

  def __repr__(self) -> str:
    return f"Tensor({self.op}, shape={self.shape})"

inverse 

inverse(eps: float = 0.0) -> 'Tensor'

1 / x (elementwise reciprocal); eps floors the divide.

Source code in aneforge/graph.py 
70
71
72
def inverse(self, eps: float = 0.0) -> "Tensor":
  """`1 / x` (elementwise reciprocal); `eps` floors the divide."""
  return Tensor(self.shape, "inverse", [self], {"eps": eps})

scaled_tanh 

scaled_tanh(alpha: float = 1.0, beta: float = 1.0) -> 'Tensor'

alpha * tanh(beta * x).

Source code in aneforge/graph.py 
78
79
80
def scaled_tanh(self, alpha: float = 1.0, beta: float = 1.0) -> "Tensor":
  """`alpha * tanh(beta * x)`."""
  return Tensor(self.shape, "scaled_tanh", [self], {"alpha": alpha, "beta": beta})

threshold 

threshold(alpha: float = 0.0) -> 'Tensor'

max(x, alpha).

Source code in aneforge/graph.py 
81
82
83
def threshold(self, alpha: float = 0.0) -> "Tensor":
  """`max(x, alpha)`."""
  return Tensor(self.shape, "threshold", [self], {"alpha": alpha})

thresholded_relu 

thresholded_relu(alpha: float = 1.0) -> 'Tensor'

x if x >= alpha else 0.

Source code in aneforge/graph.py 
84
85
86
def thresholded_relu(self, alpha: float = 1.0) -> "Tensor":
  """`x if x >= alpha else 0`."""
  return Tensor(self.shape, "thresholded_relu", [self], {"alpha": alpha})

clamped_relu 

clamped_relu(alpha: float = 0.0, beta: float = 6.0) -> 'Tensor'

min(beta, x) for x >= 0 else min(beta, alpha * x) (a leaky relu6).

Source code in aneforge/graph.py 
87
88
89
def clamped_relu(self, alpha: float = 0.0, beta: float = 6.0) -> "Tensor":
  """`min(beta, x)` for `x >= 0` else `min(beta, alpha * x)` (a leaky relu6)."""
  return Tensor(self.shape, "clamped_relu", [self], {"alpha": alpha, "beta": beta})

sigmoid_hard 

sigmoid_hard(alpha: float = 0.2, beta: float = 0.5) -> 'Tensor'

Hard sigmoid: min(max(alpha * x + beta, 0), 1).

Source code in aneforge/graph.py 
90
91
92
def sigmoid_hard(self, alpha: float = 0.2, beta: float = 0.5) -> "Tensor":
  """Hard sigmoid: `min(max(alpha * x + beta, 0), 1)`."""
  return Tensor(self.shape, "sigmoid_hard", [self], {"alpha": alpha, "beta": beta})

linear_activation 

linear_activation(alpha: float = 1.0, beta: float = 0.0) -> 'Tensor'

alpha * x + beta (scalar affine, fused as one op).

Source code in aneforge/graph.py 
93
94
95
def linear_activation(self, alpha: float = 1.0, beta: float = 0.0) -> "Tensor":
  """`alpha * x + beta` (scalar affine, fused as one op)."""
  return Tensor(self.shape, "linear_activation", [self], {"alpha": alpha, "beta": beta})

greater 

greater(o) -> 'Tensor'

Elementwise x > o -> a BOOL tensor (use as the cond of af.select).

Source code in aneforge/graph.py 
 97
 98
 99
100
101
def greater(self, o) -> "Tensor":
  """Elementwise `x > o` -> a BOOL tensor (use as the `cond` of `af.select`)."""
  if not isinstance(o, Tensor):
    raise TypeError("greater expects a graph Tensor (compare against a streamed weight via select)")
  return Tensor(_broadcast(self.shape, o.shape), "greater", [self, o])

prelu 

prelu(alpha) -> 'Tensor'

Per-channel PReLU: x if x>0 else alpha[c]*x. alpha: [C]; input rank>=3 [N,C,...].

Source code in aneforge/graph.py 
120
121
122
123
124
125
126
127
def prelu(self, alpha) -> "Tensor":
  """Per-channel PReLU: `x if x>0 else alpha[c]*x`. `alpha`: [C]; input rank>=3 [N,C,...]."""
  alpha = np.asarray(alpha)
  if len(self.shape) < 3:
    raise ValueError(f"prelu needs rank>=3 [N,C,...]; got {self.shape}")
  if alpha.shape != (self.shape[1],):
    raise ValueError(f"prelu alpha {alpha.shape} != channels {self.shape[1]}")
  return Tensor(self.shape, "prelu", [self], {"alpha": alpha})

reverse 

reverse(axes) -> 'Tensor'

Reverse along axes (native reverse).

Source code in aneforge/graph.py 
132
133
134
135
136
137
def reverse(self, axes) -> "Tensor":
  """Reverse along `axes` (native `reverse`)."""
  if not isinstance(axes, (tuple, list)):
    axes = (axes,)
  axes = tuple(a % len(self.shape) for a in axes)
  return Tensor(self.shape, "reverse", [self], {"axes": axes})

tile 

tile(reps) -> 'Tensor'

Repeat reps[i] times along each axis (native tile; factors of {2,3,4,8}).

Source code in aneforge/graph.py 
139
140
141
142
143
def tile(self, reps) -> "Tensor":
  """Repeat `reps[i]` times along each axis (native `tile`; factors of {2,3,4,8})."""
  reps = tuple(int(r) for r in reps)
  out = tuple(d * r for d, r in zip(self.shape, reps))
  return Tensor(out, "tile", [self], {"reps": reps})

reduce_log_sum_exp 

reduce_log_sum_exp(axes) -> 'Tensor'

log(sum(exp(x))) over axes (native, stable softmax denominator).

Source code in aneforge/graph.py 
145
146
147
def reduce_log_sum_exp(self, axes) -> "Tensor":
  """log(sum(exp(x))) over `axes` (native, stable softmax denominator)."""
  return self._reduce("reduce_log_sum_exp", axes)

adds 

adds(k: float) -> 'Tensor'

x + scalar as a fused scalar-add (the only fused way to inject a scalar offset; + needs two Tensors).

Source code in aneforge/graph.py 
161
162
163
def adds(self, k: float) -> "Tensor":
  """`x + scalar` as a fused scalar-add (the only fused way to inject a scalar offset; `+` needs two Tensors)."""
  return Tensor(self.shape, "adds", [self], {"k": float(k)})

bmm_weight 

bmm_weight(W) -> 'Tensor'

Batched matmul self [B,M,K] @ W [B,K,N], W a baked QUANTIZABLE weight (unlike @ with a _const, which stays fp16) -- for per-expert projections, so compress= shrinks them.

Source code in aneforge/graph.py 
179
180
181
182
183
184
185
def bmm_weight(self, W) -> "Tensor":
  """Batched matmul `self [B,M,K] @ W [B,K,N]`, `W` a baked QUANTIZABLE weight (unlike `@` with a `_const`,
  which stays fp16) -- for per-expert projections, so `compress=` shrinks them."""
  W = np.asarray(W); _check_dtype(W, "bmm_weight")
  if W.ndim != 3 or self.shape[-1] != W.shape[-2]:
    raise ValueError(f"bmm_weight: x{self.shape} @ W{W.shape} shape mismatch")
  return Tensor(self.shape[:-1] + (W.shape[-1],), "bmm_w", [self], {"wt": np.ascontiguousarray(W)})

linear 

linear(W, bias=None) -> 'Tensor'

x @ W.T (+ bias). W is [out, in] (PyTorch convention).

Source code in aneforge/graph.py 
187
188
189
190
191
192
193
194
195
def linear(self, W, bias=None) -> "Tensor":
  """`x @ W.T (+ bias)`. `W` is [out, in] (PyTorch convention)."""
  W = np.asarray(W); _check_dtype(W, "linear weight")
  if W.ndim != 2 or self.shape[-1] != W.shape[1]:
    raise ValueError(f"linear: x{self.shape} W{W.shape} shape mismatch")
  attrs: dict[str, Any] = {"wt": np.ascontiguousarray(W)}
  if bias is not None:
    attrs["bias"] = np.asarray(bias).astype(np.float32)
  return Tensor(self.shape[:-1] + (W.shape[0],), "matmul", [self], attrs)

squeeze 

squeeze(axes) -> 'Tensor'

Remove size-1 dims at axes (native squeeze).

Source code in aneforge/graph.py 
206
207
208
209
210
211
212
213
214
215
def squeeze(self, axes) -> "Tensor":
  """Remove size-1 dims at `axes` (native `squeeze`)."""
  if not isinstance(axes, (tuple, list)):
    axes = (axes,)
  axes = tuple(a % len(self.shape) for a in axes)
  for a in axes:
    if self.shape[a] != 1:
      raise ValueError(f"squeeze: axis {a} has size {self.shape[a]} != 1")
  out = tuple(d for i, d in enumerate(self.shape) if i not in axes)
  return Tensor(out, "squeeze", [self], {"axes": axes})

expand_dims 

expand_dims(axes) -> 'Tensor'

Insert size-1 dims at axes (native expand_dims); axes index the OUTPUT rank.

Source code in aneforge/graph.py 
217
218
219
220
221
222
223
224
225
226
def expand_dims(self, axes) -> "Tensor":
  """Insert size-1 dims at `axes` (native `expand_dims`); axes index the OUTPUT rank."""
  if not isinstance(axes, (tuple, list)):
    axes = (axes,)
  out_rank = len(self.shape) + len(axes)
  axes = sorted(a % out_rank for a in axes)
  out, src = [], iter(self.shape)
  for i in range(out_rank):
    out.append(1 if i in axes else next(src))
  return Tensor(tuple(out), "expand_dims", [self], {"axes": tuple(axes)})

flatten2d 

flatten2d(axis: int = 1) -> 'Tensor'

Collapse to 2-D about axis: [:axis] -> rows, [axis:] -> cols (native flatten2d).

Source code in aneforge/graph.py 
228
229
230
231
232
233
def flatten2d(self, axis: int = 1) -> "Tensor":
  """Collapse to 2-D about `axis`: `[:axis]` -> rows, `[axis:]` -> cols (native `flatten2d`)."""
  ax = axis % len(self.shape)
  rows = int(np.prod(self.shape[:ax])) if ax > 0 else 1
  cols = int(np.prod(self.shape[ax:]))
  return Tensor((rows, cols), "flatten2d", [self], {"axis": ax})

slice_by_size 

slice_by_size(begin, size) -> 'Tensor'

Static per-axis slice x[begin[i]:begin[i]+size[i]] (native slice_by_size).

Source code in aneforge/graph.py 
235
236
237
238
239
240
241
242
243
def slice_by_size(self, begin, size) -> "Tensor":
  """Static per-axis slice `x[begin[i]:begin[i]+size[i]]` (native `slice_by_size`)."""
  begin = [int(b) for b in begin]; size = [int(s) for s in size]
  if len(begin) != len(self.shape) or len(size) != len(self.shape):
    raise ValueError(f"slice_by_size: begin/size must have rank {len(self.shape)}")
  for i, (b, s) in enumerate(zip(begin, size)):
    if b < 0 or s <= 0 or b + s > self.shape[i]:
      raise ValueError(f"slice_by_size: axis {i} window [{b}:{b+s}] out of range for {self.shape[i]}")
  return Tensor(tuple(size), "slice_by_size", [self], {"begin": begin, "size": size})

cumsum 

cumsum(axis: int = -1) -> 'Tensor'

Cumulative sum along the last axis as x @ triu_ones (no native cumsum).

Source code in aneforge/graph.py 
258
259
260
261
262
263
264
265
266
def cumsum(self, axis: int = -1) -> "Tensor":
  """Cumulative sum along the last axis as `x @ triu_ones` (no native cumsum)."""
  ax = axis % len(self.shape)
  if ax != len(self.shape) - 1:
    raise ValueError(f"cumsum: only the last axis is supported (got axis={axis}, "
            f"rank {len(self.shape)}); transpose the axis to last first")
  N = self.shape[-1]
  W = np.tril(np.ones((N, N), dtype=np.float32)).astype(np.float16)  # linear(W)=x@W.T=x@triu
  return self.linear(W, None)

l1_norm 

l1_norm(axes) -> 'Tensor'

sum(|x|, axes) (keepdims) via the native reduce_l1_norm op.

Source code in aneforge/graph.py 
267
268
269
def l1_norm(self, axes) -> "Tensor":
  """`sum(|x|, axes)` (keepdims) via the native `reduce_l1_norm` op."""
  return self._reduce("reduce_l1_norm", axes)

log_sum 

log_sum(axes) -> 'Tensor'

log(sum(x, axes)) (keepdims). Expects a positive input.

Source code in aneforge/graph.py 
270
271
272
def log_sum(self, axes) -> "Tensor":
  """`log(sum(x, axes))` (keepdims). Expects a positive input."""
  return self._reduce("reduce_log_sum", axes)

sum_square 

sum_square(axes) -> 'Tensor'

sum(x**2, axes) (keepdims) via the native reduce_sum_square op.

Source code in aneforge/graph.py 
273
274
275
def sum_square(self, axes) -> "Tensor":
  """`sum(x**2, axes)` (keepdims) via the native `reduce_sum_square` op."""
  return self._reduce("reduce_sum_square", axes)

l2_norm 

l2_norm(axis: int = -1, eps: float = 1e-12) -> 'Tensor'

Per-axis L2-normalize x / sqrt(sum(x**2, axis) + eps) (built per-axis; MIL l2_norm is all-dims).

Source code in aneforge/graph.py 
280
281
282
283
def l2_norm(self, axis: int = -1, eps: float = 1e-12) -> "Tensor":
  """Per-axis L2-normalize `x / sqrt(sum(x**2, axis) + eps)` (built per-axis; MIL `l2_norm` is all-dims)."""
  ax = axis % len(self.shape)
  return Tensor(self.shape, "l2_norm", [self], {"axis": ax, "eps": float(eps)})

argmax 

argmax(axis: int = -1) -> 'Tensor'

Argmax along axis (keepdims). GlobalArgMinMax bridge (cut); 2D [C,W] only, indices fp16-encoded.

Source code in aneforge/graph.py 
285
286
287
288
289
290
291
def argmax(self, axis: int = -1) -> "Tensor":
  """Argmax along `axis` (keepdims). GlobalArgMinMax bridge (cut); 2D [C,W] only, indices fp16-encoded."""
  if len(self.shape) != 2:
    raise ValueError(f"argmax: only 2D [C,W] inputs are supported; got {self.shape}")
  ax = axis % 2
  out = tuple(1 if i == ax else d for i, d in enumerate(self.shape))
  return Tensor(out, "argmax", [self], {"axis": ax})

rms_norm 

rms_norm(gamma, eps: float = 1e-05) -> 'Tensor'

RMSNorm over the last dim. gamma: a [D] array (baked) or a [1,D] Tensor (trainable).

Source code in aneforge/graph.py 
293
294
295
296
297
298
299
300
301
def rms_norm(self, gamma, eps: float = 1e-5) -> "Tensor":
  """RMSNorm over the last dim. `gamma`: a [D] array (baked) or a [1,D] Tensor (trainable)."""
  if isinstance(gamma, Tensor):
    xn = self.rms_norm(np.ones(self.shape[-1], np.float32), eps)
    return xn * gamma
  gamma = np.asarray(gamma); _check_dtype(gamma, "rms_norm gamma")
  if len(self.shape) != 2 or gamma.shape != (self.shape[-1],):
    raise ValueError(f"rms_norm expects 2D [M,D] with gamma [D]; got {self.shape}, {gamma.shape}")
  return Tensor(self.shape, "rms_norm", [self], {"gamma": gamma, "eps": float(eps)})

layer_norm 

layer_norm(gamma, beta, eps: float = 1e-05) -> 'Tensor'

LayerNorm over the last dim (2D [M,D]). gamma/beta: [D] arrays (baked) or [1,D] Tensors (trainable; pass both).

Source code in aneforge/graph.py 
303
304
305
306
307
308
309
310
311
312
313
314
315
def layer_norm(self, gamma, beta, eps: float = 1e-5) -> "Tensor":
  """LayerNorm over the last dim (2D [M,D]). `gamma`/`beta`: [D] arrays (baked) or [1,D] Tensors (trainable; pass both)."""
  if isinstance(gamma, Tensor) or isinstance(beta, Tensor):
    if not (isinstance(gamma, Tensor) and isinstance(beta, Tensor)):
      raise TypeError("layer_norm: a trainable affine needs both gamma and beta as Tensors")
    D = self.shape[-1]
    xn = self.layer_norm(np.ones(D, np.float32), np.zeros(D, np.float32), eps)
    return xn * gamma + beta
  gamma = np.asarray(gamma); beta = np.asarray(beta)
  _check_dtype(gamma, "layer_norm gamma"); _check_dtype(beta, "layer_norm beta")
  if len(self.shape) != 2 or gamma.shape != (self.shape[-1],):
    raise ValueError(f"layer_norm expects 2D [M,D] with gamma/beta [D]; got {self.shape}, {gamma.shape}")
  return Tensor(self.shape, "layer_norm", [self], {"gamma": gamma, "beta": beta, "eps": float(eps)})

channel_layer_norm 

channel_layer_norm(gamma, beta, eps: float = 1e-05) -> 'Tensor'

LayerNorm over the CHANNEL axis of [N,C,1,S] (ANE transformer layout; no transpose). gamma/beta: [C].

Source code in aneforge/graph.py 
317
318
319
320
321
322
323
324
325
def channel_layer_norm(self, gamma, beta, eps: float = 1e-5) -> "Tensor":
  """LayerNorm over the CHANNEL axis of [N,C,1,S] (ANE transformer layout; no transpose). `gamma`/`beta`: [C]."""
  gamma = np.asarray(gamma); beta = np.asarray(beta)
  _check_dtype(gamma, "channel_layer_norm gamma"); _check_dtype(beta, "channel_layer_norm beta")
  if (len(self.shape) != 4 or self.shape[2] != 1
      or gamma.shape != (self.shape[1],) or beta.shape != (self.shape[1],)):
    raise ValueError(f"channel_layer_norm expects [N,C,1,S] with gamma/beta [C]; "
            f"got {self.shape}, {gamma.shape}")
  return Tensor(self.shape, "channel_layer_norm", [self], {"gamma": gamma, "beta": beta, "eps": float(eps)})

group_norm 

group_norm(gamma, beta, num_groups: int, eps: float = 1e-05) -> 'Tensor'

GroupNorm over [1,C,H,W]. gamma/beta: [C] arrays (baked) or [1,C,1,1] Tensors (trainable; pass both).

Source code in aneforge/graph.py 
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
def group_norm(self, gamma, beta, num_groups: int, eps: float = 1e-5) -> "Tensor":
  """GroupNorm over [1,C,H,W]. `gamma`/`beta`: [C] arrays (baked) or [1,C,1,1] Tensors (trainable; pass both)."""
  if isinstance(gamma, Tensor) or isinstance(beta, Tensor):
    if not (isinstance(gamma, Tensor) and isinstance(beta, Tensor)):
      raise TypeError("group_norm: a trainable affine needs both gamma and beta as Tensors")
    C = self.shape[1]
    xn = self.group_norm(np.ones(C, np.float32), np.zeros(C, np.float32), num_groups, eps)
    return xn * gamma + beta
  gamma = np.asarray(gamma); beta = np.asarray(beta)
  _check_dtype(gamma, "group_norm gamma"); _check_dtype(beta, "group_norm beta")
  if len(self.shape) != 4 or self.shape[0] != 1 or self.shape[1] % num_groups:
    raise ValueError(f"group_norm expects [1,C,H,W] with C%groups==0; got {self.shape}, G={num_groups}")
  # tiled lowering reduces over max(C/groups, H*W) vs the 65536 per-axis cap (finding_sd15)
  _, C, H, W = self.shape
  axis = max(C // num_groups, H * W)
  if axis > 65536:
    raise ValueError(
      f"group_norm: largest tiled axis max(C/groups, H*W) = {axis} exceeds "
      f"the ANE per-axis bound (65536) for {self.shape} with groups={num_groups}; "
      f"reduce the feature map or channels, or use more groups.")
  return Tensor(self.shape, "group_norm", [self],
         {"gamma": gamma, "beta": beta, "groups": num_groups, "eps": float(eps)})

avg_pool 

avg_pool(k: int, stride: int | None = None, pad: int = 0, ceil_mode: bool = False, exclude_pad: bool = False) -> 'Tensor'

Average pool. exclude_pad divides by the valid (non-pad) cell count (ONNX count_include_pad=0); default includes pad cells (divide by full kernel area).

Source code in aneforge/graph.py 
363
364
365
366
def avg_pool(self, k: int, stride: int | None = None, pad: int = 0,
             ceil_mode: bool = False, exclude_pad: bool = False) -> "Tensor":
  """Average pool. `exclude_pad` divides by the valid (non-pad) cell count (ONNX `count_include_pad=0`); default includes pad cells (divide by full kernel area)."""
  return self._pool("avg_pool", k, stride, pad, ceil_mode, exclude_pad)

upsample 

upsample(scale: int = 2) -> 'Tensor'

Nearest-neighbour upsample [N,C,H,W] -> [N,C,scaleH,scaleW].

Source code in aneforge/graph.py 
368
369
370
371
def upsample(self, scale: int = 2) -> "Tensor":
  """Nearest-neighbour upsample [N,C,H,W] -> [N,C,scale*H,scale*W]."""
  N, C, H, W = self.shape
  return Tensor((N, C, H * scale, W * scale), "upsample", [self], {"scale": scale})

input 

input(shape: Sequence[int], dtype: str = 'fp16', max_abs: 'float | None' = None) -> Tensor

A graph input placeholder (fed to the Model in creation order). dtype: "fp16" or "uint8" (raw bytes, dequantised in-graph; see af.image_input).

max_abs declares a bound on |value| at runtime, which lets the optimizer keep a lossy variant whose encoding range provably covers this graph (see af.compile(opt=1) and #155). It is a promise, not a clamp: nothing enforces it at dispatch, and feeding larger values makes the bound wrong. Left undeclared the bound is unknown, and the optimizer stays fail-closed.

Source code in aneforge/graph.py 
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
def input(shape: Sequence[int], dtype: str = "fp16", max_abs: "float | None" = None) -> Tensor:
  """A graph input placeholder (fed to the Model in creation order). `dtype`: "fp16" or "uint8" (raw bytes, dequantised in-graph; see `af.image_input`).

  `max_abs` declares a bound on `|value|` at runtime, which lets the optimizer keep a lossy variant
  whose encoding range provably covers this graph (see `af.compile(opt=1)` and #155). It is a promise,
  not a clamp: nothing enforces it at dispatch, and feeding larger values makes the bound wrong. Left
  undeclared the bound is unknown, and the optimizer stays fail-closed."""
  if dtype not in ("fp16", "uint8"):
    raise ValueError(f"input: dtype must be 'fp16' or 'uint8'; got {dtype!r}")
  t = Tensor(tuple(shape), "input")
  t.attrs["idx"] = _input_counter[0]
  if max_abs is not None:
    max_abs = float(max_abs)
    if not (max_abs >= 0.0):          # rejects negatives and nan
      raise ValueError(f"input: max_abs must be a non-negative bound on |value|; got {max_abs!r}")
    t.attrs["max_abs"] = max_abs
  if dtype != "fp16":
    t.attrs["dtype"] = dtype
  _input_counter[0] += 1
  return t

image_input 

image_input(shape: Sequence[int], scale: float = 1.0 / 255.0, bias=0.0) -> Tensor

A uint8 image input dequantised to fp16 on-engine (cast -> mul(scale) -> add(bias)). scale/bias are scalar or length-C (NCHW per-channel, broadcast [1,C,1,1]).

Source code in aneforge/graph.py 
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
def image_input(shape: Sequence[int], scale: float = 1.0 / 255.0, bias=0.0) -> Tensor:
  """A uint8 image input dequantised to fp16 on-engine (`cast -> mul(scale) -> add(bias)`). `scale`/`bias` are scalar or length-C (NCHW per-channel, broadcast [1,C,1,1])."""
  shape = tuple(int(d) for d in shape)
  x = input(shape, dtype="uint8")
  y = Tensor(shape, "cast", [x], {"dtype": "fp16"})

  def _coef(v, what: str):
    arr = np.asarray(v, dtype=np.float32)
    if arr.ndim == 0:
      return float(arr), None
    if len(shape) != 4 or arr.shape != (shape[1],):
      raise ValueError(f"image_input: per-channel {what} must be length-C={shape[1]} "
              f"for an NCHW image; got {arr.shape} with shape {shape}")
    return None, arr.reshape(1, shape[1], 1, 1).astype(np.float16)

  sc_s, sc_v = _coef(scale, "scale")
  if sc_v is not None:                                    # per-channel scale: broadcast mul
    y = y * Tensor(sc_v.shape, "const_array", [], {"value": sc_v})
  elif sc_s != 1.0:
    y = y * sc_s
  bi_s, bi_v = _coef(bias, "bias")
  if bi_v is not None:                                    # per-channel bias: broadcast add
    y = y + Tensor(bi_v.shape, "const_array", [], {"value": bi_v})
  elif bi_s != 0.0:
    assert bi_s is not None
    y = y.adds(bi_s)
  return y

conv 

conv(x: Tensor, weight, stride: int = 1, pad: 'int | tuple[int, int, int, int]' = 0, dilation: int = 1, groups: int = 1, bias=None) -> Tensor

2D conv. x: [N,Cin,H,W]; weight: [Cout, Cin/groups, kH, kW]; bias: [Cout]. pad is a scalar (symmetric) or a (top, bottom, left, right) tuple.

Source code in aneforge/graph.py 
465
466
467
468
469
470
471
472
473
474
475
476
477
478
def conv(x: Tensor, weight, stride: int = 1, pad: "int | tuple[int, int, int, int]" = 0, dilation: int = 1,
    groups: int = 1, bias=None) -> Tensor:
  """2D conv. `x`: [N,Cin,H,W]; `weight`: [Cout, Cin/groups, kH, kW]; `bias`: [Cout]. `pad` is a scalar (symmetric) or a `(top, bottom, left, right)` tuple."""
  weight = np.asarray(weight); _check_dtype(weight, "conv weight")
  if len(x.shape) != 4:
    raise ValueError(f"conv expects 4D input [N,Cin,H,W], got {x.shape}")
  N, Cin, H, W = x.shape
  Cout, _, kH, kW = weight.shape
  _check_kw(kW, "conv", weight.shape)
  pt, pb, pl, pr = (pad, pad, pad, pad) if isinstance(pad, int) else pad
  Hout = (H + pt + pb - dilation * (kH - 1) - 1) // stride + 1
  Wout = (W + pl + pr - dilation * (kW - 1) - 1) // stride + 1
  return Tensor((N, Cout, Hout, Wout), "conv", [x],
         _conv_attrs(weight, stride, pad, dilation, groups, bias))

dynamic_conv 

dynamic_conv(x: Tensor, weight: Tensor, stride: int = 1, pad: int = 0, dilation: int = 1, groups: int = 1) -> Tensor

2D conv with a runtime-tensor weight (native dynamic-kernel path; hypernetworks/per-sample kernels). x: [1,Cin,H,W]; weight Tensor [Cout,Cin/g,kH,kW]; batch must be 1.

Source code in aneforge/graph.py 
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
def dynamic_conv(x: Tensor, weight: Tensor, stride: int = 1, pad: int = 0,
        dilation: int = 1, groups: int = 1) -> Tensor:
  """2D conv with a runtime-tensor weight (native dynamic-kernel path; hypernetworks/per-sample kernels). `x`: [1,Cin,H,W]; `weight` Tensor [Cout,Cin/g,kH,kW]; batch must be 1."""
  if not isinstance(weight, Tensor):
    raise TypeError("dynamic_conv: weight must be a Tensor (a graph value); for a constant "
            "weight use af.conv")
  if len(x.shape) != 4 or len(weight.shape) != 4:
    raise ValueError(f"dynamic_conv: x and weight must be 4D [N,Cin,H,W]/[Cout,Cin/g,kH,kW]; "
            f"got {x.shape}, {weight.shape}")
  N, Cin, H, W = x.shape
  Cout, Cin_g, kH, kW = weight.shape
  if N != 1:
    raise ValueError(
      f"dynamic_conv requires batch N=1 (got N={N}): a dynamic-weight conv with batch>=2 "
      f"is unsupported on the ANE dynamic-kernel path. Use af.conv / conv2d for batched convolution.")
  if Cin_g * groups != Cin:
    raise ValueError(f"dynamic_conv: weight Cin/groups={Cin_g} x groups={groups} != input Cin={Cin}")
  _check_kw(kW, "dynamic_conv", None)
  Hout = (H + 2 * pad - dilation * (kH - 1) - 1) // stride + 1
  Wout = (W + 2 * pad - dilation * (kW - 1) - 1) // stride + 1
  return Tensor((N, Cout, Hout, Wout), "dynamic_conv", [x, weight],
         {"stride": stride, "pad": pad, "dilation": dilation, "groups": groups})

conv_transpose 

conv_transpose(x: Tensor, weight, stride: int = 1, pad: int = 0, dilation: int = 1, groups: int = 1, bias=None) -> Tensor

2D transposed conv (deconv). x: [N,Cin,H,W]; weight: [Cin,Cout,kH,kW] (PyTorch layout); bias: [Cout].

Source code in aneforge/graph.py 
505
506
507
508
509
510
511
512
513
514
515
516
517
def conv_transpose(x: Tensor, weight, stride: int = 1, pad: int = 0, dilation: int = 1,
         groups: int = 1, bias=None) -> Tensor:
  """2D transposed conv (deconv). `x`: [N,Cin,H,W]; `weight`: [Cin,Cout,kH,kW] (PyTorch layout); `bias`: [Cout]."""
  weight = np.asarray(weight); _check_dtype(weight, "conv_transpose weight")
  if len(x.shape) != 4:
    raise ValueError(f"conv_transpose expects 4D [N,Cin,H,W], got {x.shape}")
  N, Cin, H, W = x.shape
  _, Cout, kH, kW = weight.shape
  _check_kw(kW, "conv_transpose", weight.shape)
  Hout = (H - 1) * stride - 2 * pad + dilation * (kH - 1) + 1
  Wout = (W - 1) * stride - 2 * pad + dilation * (kW - 1) + 1
  return Tensor((N, Cout, Hout, Wout), "conv_transpose", [x],
         _conv_attrs(weight, stride, pad, dilation, groups, bias))

batch_norm 

batch_norm(x: Tensor, gamma, beta, mean, var, eps: float = 1e-05) -> Tensor

Inference BatchNorm over [1,C,...] from running mean/var. gamma/beta/mean/var: [C].

Source code in aneforge/graph.py 
520
521
522
523
524
525
526
527
528
def batch_norm(x: Tensor, gamma, beta, mean, var, eps: float = 1e-5) -> Tensor:
  """Inference BatchNorm over [1,C,...] from running `mean`/`var`. gamma/beta/mean/var: [C]."""
  g, b, m, v = (np.asarray(a) for a in (gamma, beta, mean, var))
  for a, nm in ((g, "gamma"), (b, "beta"), (m, "mean"), (v, "var")):
    _check_dtype(a, f"batch_norm {nm}")
  if len(x.shape) < 3 or g.shape != (x.shape[1],):
    raise ValueError(f"batch_norm expects rank>=3 [1,C,...] with params [C]; got {x.shape}, {g.shape}")
  return Tensor(x.shape, "batch_norm", [x],
         {"gamma": g, "beta": b, "mean": m, "var": v, "eps": float(eps)})

maximum 

maximum(a: Tensor, b: Tensor) -> Tensor

Elementwise max of two graph tensors.

Source code in aneforge/graph.py 
531
532
533
def maximum(a: Tensor, b: Tensor) -> Tensor:
  """Elementwise max of two graph tensors."""
  return _binary(a, b, "maximum")

minimum 

minimum(a: Tensor, b: Tensor) -> Tensor

Elementwise min of two graph tensors.

Source code in aneforge/graph.py 
536
537
538
def minimum(a: Tensor, b: Tensor) -> Tensor:
  """Elementwise min of two graph tensors."""
  return _binary(a, b, "minimum")

concat 

concat(tensors: Sequence[Tensor], axis: int = 1) -> Tensor

Concatenate tensors along axis (e.g. UNet skip connections).

Source code in aneforge/graph.py 
541
542
543
544
545
546
547
def concat(tensors: Sequence[Tensor], axis: int = 1) -> Tensor:
  """Concatenate tensors along `axis` (e.g. UNet skip connections)."""
  tensors = list(tensors)
  ax = axis % len(tensors[0].shape)
  out = list(tensors[0].shape)
  out[ax] = sum(t.shape[ax] for t in tensors)
  return Tensor(tuple(out), "concat", tensors, {"axis": ax})

gather 

gather(x: Tensor, indices, axis: int = 0) -> Tensor

Gather slices along axis by STATIC integer indices via slice_by_size + concat (no native gather).

Source code in aneforge/graph.py 
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
def gather(x: Tensor, indices, axis: int = 0) -> Tensor:
  """Gather slices along `axis` by STATIC integer `indices` via `slice_by_size` + `concat` (no native gather)."""
  idx = [int(i) for i in indices]
  rank = len(x.shape)
  ax = axis % rank
  # last-axis gather hits the A13/A14 x16 crop-DMA bug; transpose the axis off last
  if ax == rank - 1:
    if rank == 1:
      return gather(x.reshape(x.shape[0], 1), idx, axis=0).reshape(len(idx))
    perm = list(range(rank)); perm[ax], perm[-2] = perm[-2], perm[ax]
    return gather(x.transpose(perm), idx, axis=rank - 2).transpose(perm)
  rows = []
  for i in idx:
    if not (0 <= i < x.shape[ax]):
      raise ValueError(f"gather: index {i} out of range for axis {ax} (size {x.shape[ax]})")
    begin = [0] * rank; begin[ax] = i
    size = list(x.shape); size[ax] = 1
    rows.append(x.slice_by_size(begin, size))
  return concat(rows, axis=ax)

stack 

stack(tensors: Sequence[Tensor], axis: int = 0) -> Tensor

Stack equal-shaped tensors along a NEW axis (native stack), N inserted at axis.

Source code in aneforge/graph.py 
571
572
573
574
575
576
577
578
579
580
581
582
def stack(tensors: Sequence[Tensor], axis: int = 0) -> Tensor:
  """Stack equal-shaped tensors along a NEW axis (native `stack`), N inserted at `axis`."""
  tensors = list(tensors)
  if not tensors:
    raise ValueError("stack: empty tensor list")
  base = tensors[0].shape
  for t in tensors:
    if t.shape != base:
      raise ValueError(f"stack: all tensors must share a shape; got {base} and {t.shape}")
  ax = axis % (len(base) + 1)
  out = base[:ax] + (len(tensors),) + base[ax:]
  return Tensor(out, "stack", tensors, {"axis": ax})

split 

split(x: Tensor, num_splits: int, axis: int = 0) -> list[Tensor]

Split x into num_splits equal parts along axis (native split); axis size must divide evenly.

Source code in aneforge/graph.py 
585
586
587
588
589
590
591
592
593
def split(x: Tensor, num_splits: int, axis: int = 0) -> list[Tensor]:
  """Split `x` into `num_splits` equal parts along `axis` (native `split`); axis size must divide evenly."""
  ax = axis % len(x.shape)
  if x.shape[ax] % num_splits:
    raise ValueError(f"split: axis {ax} size {x.shape[ax]} not divisible by {num_splits}")
  part = x.shape[ax] // num_splits
  out_shape = x.shape[:ax] + (part,) + x.shape[ax + 1:]
  return [Tensor(out_shape, "split", [x], {"axis": ax, "num_splits": num_splits, "which": i})
      for i in range(num_splits)]

select 

select(cond: Tensor, a: Tensor, b: Tensor) -> Tensor

Elementwise cond ? a : b (native select). cond is a BOOL tensor; a/b are fp16.

Source code in aneforge/graph.py 
596
597
598
599
600
def select(cond: Tensor, a: Tensor, b: Tensor) -> Tensor:
  """Elementwise `cond ? a : b` (native `select`). `cond` is a BOOL tensor; `a`/`b` are fp16."""
  if not (isinstance(cond, Tensor) and isinstance(a, Tensor) and isinstance(b, Tensor)):
    raise TypeError("select expects three graph Tensors (cond, a, b)")
  return Tensor(_broadcast(a.shape, b.shape), "select", [cond, a, b])

instance_norm 

instance_norm(x: Tensor, gamma, beta, eps: float = 1e-05) -> Tensor

InstanceNorm over [N,C,H,W] (per-(N,C)-slice spatial norm + per-channel affine). gamma/beta: [C].

Source code in aneforge/graph.py 
606
607
608
609
610
611
612
def instance_norm(x: Tensor, gamma, beta, eps: float = 1e-5) -> Tensor:
  """InstanceNorm over [N,C,H,W] (per-(N,C)-slice spatial norm + per-channel affine). `gamma`/`beta`: [C]."""
  g, b = np.asarray(gamma), np.asarray(beta)
  _check_dtype(g, "instance_norm gamma"); _check_dtype(b, "instance_norm beta")
  if len(x.shape) != 4 or g.shape != (x.shape[1],) or b.shape != (x.shape[1],):
    raise ValueError(f"instance_norm expects [N,C,H,W] with gamma/beta [C]; got {x.shape}, {g.shape}")
  return Tensor(x.shape, "instance_norm", [x], {"gamma": g, "beta": b, "eps": float(eps)})

local_response_norm 

local_response_norm(x: Tensor, size: int = 5, alpha: float = 0.0001, beta: float = 0.75, k: float = 1.0) -> Tensor

Cross-channel LRN over [N,C,H,W] via the fused native local_response_norm op (no cut; distinct from the af.lrn bridge).

Source code in aneforge/graph.py 
615
616
617
618
619
620
621
def local_response_norm(x: Tensor, size: int = 5, alpha: float = 1e-4,
            beta: float = 0.75, k: float = 1.0) -> Tensor:
  """Cross-channel LRN over [N,C,H,W] via the fused native `local_response_norm` op (no cut; distinct from the `af.lrn` bridge)."""
  if len(x.shape) != 4:
    raise ValueError(f"local_response_norm expects [N,C,H,W]; got {x.shape}")
  return Tensor(x.shape, "local_response_norm", [x],
         {"size": int(size), "alpha": float(alpha), "beta": float(beta), "k": float(k)})

einsum_native 

einsum_native(equation: str, a: Tensor, b) -> Tensor

Restricted batched contraction via the native einsum op (distinct from af.einsum). Only 'nchw,nwhu->nchu' is reachable: a=[N,C,H,W], b=[N,W,H,U] -> [N,C,H,U].

Source code in aneforge/graph.py 
624
625
626
627
628
629
630
631
632
633
634
635
636
def einsum_native(equation: str, a: Tensor, b) -> Tensor:
  """Restricted batched contraction via the native `einsum` op (distinct from `af.einsum`). Only `'nchw,nwhu->nchu'` is reachable: a=[N,C,H,W], b=[N,W,H,U] -> [N,C,H,U]."""
  if equation.replace(" ", "") != "nchw,nwhu->nchu":
    raise NotImplementedError(
      f"einsum: only 'nchw,nwhu->nchu' is verified reachable on the ANE; got {equation!r}")
  b = np.asarray(b); _check_dtype(b, "einsum operand b")
  if len(a.shape) != 4 or b.ndim != 4:
    raise ValueError(f"einsum nchw,nwhu->nchu expects rank-4 a and b; got {a.shape}, {b.shape}")
  N, C, H, W = a.shape
  Nb, Wb, Hb, U = b.shape
  if (Nb, Wb, Hb) != (N, W, H):
    raise ValueError(f"einsum: b must be [N,W,H,U]=[{N},{W},{H},U]; got {b.shape}")
  return Tensor((N, C, H, U), "einsum", [a], {"b": b, "equation": "nchw,nwhu->nchu"})

space_to_depth 

space_to_depth(x: Tensor, block_size: int = 2) -> Tensor

Space-to-depth (native space_to_depth): [N,C,H,W] -> [N,C*bs*bs,H/bs,W/bs]. Fused (no cut).

Source code in aneforge/graph.py 
639
640
641
642
643
644
645
646
647
def space_to_depth(x: Tensor, block_size: int = 2) -> Tensor:
  """Space-to-depth (native `space_to_depth`): `[N,C,H,W] -> [N,C*bs*bs,H/bs,W/bs]`. Fused (no cut)."""
  if len(x.shape) != 4:
    raise ValueError(f"space_to_depth expects 4D [N,C,H,W], got {x.shape}")
  N, C, H, W = x.shape
  bs = int(block_size)
  if H % bs or W % bs:
    raise ValueError(f"space_to_depth: H={H},W={W} not divisible by block_size={bs}")
  return Tensor((N, C * bs * bs, H // bs, W // bs), "space_to_depth", [x], {"block_size": bs})

depth_to_space 

depth_to_space(x: Tensor, block_size: int = 2) -> Tensor

Depth-to-space (native depth_to_space): [N,C*bs*bs,H,W] -> [N,C,H*bs,W*bs]. Fused (no cut).

Source code in aneforge/graph.py 
650
651
652
653
654
655
656
657
658
def depth_to_space(x: Tensor, block_size: int = 2) -> Tensor:
  """Depth-to-space (native `depth_to_space`): `[N,C*bs*bs,H,W] -> [N,C,H*bs,W*bs]`. Fused (no cut)."""
  if len(x.shape) != 4:
    raise ValueError(f"depth_to_space expects 4D [N,C,H,W], got {x.shape}")
  N, C2, H, W = x.shape
  bs = int(block_size)
  if C2 % (bs * bs):
    raise ValueError(f"depth_to_space: channels {C2} not divisible by block_size^2={bs * bs}")
  return Tensor((N, C2 // (bs * bs), H * bs, W * bs), "depth_to_space", [x], {"block_size": bs})

crop 

crop(x: Tensor, top: int, bottom: int, left: int, right: int) -> Tensor

Spatial crop of [N,C,H,W]: drop top/bottom rows, left/right cols (native crop, no cut).

Source code in aneforge/graph.py 
661
662
663
664
665
666
667
668
669
670
def crop(x: Tensor, top: int, bottom: int, left: int, right: int) -> Tensor:
  """Spatial crop of [N,C,H,W]: drop `top`/`bottom` rows, `left`/`right` cols (native `crop`, no cut)."""
  if len(x.shape) != 4:
    raise ValueError(f"crop expects 4D [N,C,H,W], got {x.shape}")
  N, C, H, W = x.shape
  Hout, Wout = H - top - bottom, W - left - right
  if Hout <= 0 or Wout <= 0:
    raise ValueError(f"crop: result {Hout}x{Wout} is empty for input {H}x{W}")
  return Tensor((N, C, Hout, Wout), "crop", [x],
         {"crop_h": (int(top), int(bottom)), "crop_w": (int(left), int(right))})

resize_nearest_neighbor 

resize_nearest_neighbor(x: Tensor, target_h: int, target_w: int) -> Tensor

Nearest-neighbour resize of [N,C,H,W] to (target_h, target_w) (native, no cut).

Source code in aneforge/graph.py 
673
674
675
676
677
678
679
def resize_nearest_neighbor(x: Tensor, target_h: int, target_w: int) -> Tensor:
  """Nearest-neighbour resize of [N,C,H,W] to `(target_h, target_w)` (native, no cut)."""
  if len(x.shape) != 4:
    raise ValueError(f"resize_nearest_neighbor expects 4D [N,C,H,W], got {x.shape}")
  N, C, _, _ = x.shape
  return Tensor((N, C, int(target_h), int(target_w)), "resize_nearest_neighbor", [x],
         {"target_h": int(target_h), "target_w": int(target_w)})

resize_bilinear 

resize_bilinear(x: Tensor, target_h: int, target_w: int, align_corners: bool = False) -> Tensor

Bilinear resize of [N,C,H,W] to (target_h, target_w) (native, no cut). Half-pixel sampling by default.

Source code in aneforge/graph.py 
682
683
684
685
686
687
688
689
690
def resize_bilinear(x: Tensor, target_h: int, target_w: int,
          align_corners: bool = False) -> Tensor:
  """Bilinear resize of [N,C,H,W] to `(target_h, target_w)` (native, no cut). Half-pixel sampling by default."""
  if len(x.shape) != 4:
    raise ValueError(f"resize_bilinear expects 4D [N,C,H,W], got {x.shape}")
  N, C, _, _ = x.shape
  return Tensor((N, C, int(target_h), int(target_w)), "resize_bilinear", [x],
         {"target_h": int(target_h), "target_w": int(target_w),
         "align_corners": bool(align_corners)})

upsample_bilinear 

upsample_bilinear(x: Tensor, scale: int = 2, align_corners: bool = False) -> Tensor

Bilinear upsample of [N,C,H,W] by integer scale (native, no cut). Half-pixel sampling by default.

Source code in aneforge/graph.py 
693
694
695
696
697
698
699
700
def upsample_bilinear(x: Tensor, scale: int = 2, align_corners: bool = False) -> Tensor:
  """Bilinear upsample of [N,C,H,W] by integer `scale` (native, no cut). Half-pixel sampling by default."""
  if len(x.shape) != 4:
    raise ValueError(f"upsample_bilinear expects 4D [N,C,H,W], got {x.shape}")
  N, C, H, W = x.shape
  s = int(scale)
  return Tensor((N, C, H * s, W * s), "upsample_bilinear", [x],
         {"scale": s, "align_corners": bool(align_corners)})

affine 

affine(x: Tensor, transform, output_h: int, output_w: int, align_corners: bool = False) -> Tensor

2-D affine warp of [N,C,H,W] to (output_h, output_w) (native affine, no cut). transform: [N,6] [a0,a1,a2,b0,b1,b2] in normalized [-1,1] coords.

Source code in aneforge/graph.py 
703
704
705
706
707
708
709
710
711
712
713
714
def affine(x: Tensor, transform, output_h: int, output_w: int,
     align_corners: bool = False) -> Tensor:
  """2-D affine warp of [N,C,H,W] to `(output_h, output_w)` (native `affine`, no cut). `transform`: [N,6] `[a0,a1,a2,b0,b1,b2]` in normalized [-1,1] coords."""
  if len(x.shape) != 4:
    raise ValueError(f"affine expects 4D [N,C,H,W], got {x.shape}")
  T = np.asarray(transform); _check_dtype(T, "affine transform")
  if T.ndim != 2 or T.shape[1] != 6:
    raise ValueError(f"affine: transform must be [N,6] (or [1,6]); got {T.shape}")
  N, C, _, _ = x.shape
  return Tensor((N, C, int(output_h), int(output_w)), "affine", [x],
         {"transform": T, "output_h": int(output_h), "output_w": int(output_w),
         "align_corners": bool(align_corners)})

pixel_shuffle 

pixel_shuffle(x: Tensor, r: int) -> Tensor

Depth-to-space upscale (PyTorch nn.PixelShuffle): [N,C*r*r,H,W] -> [N,C,H*r,W*r]. Fused (no cut).

Source code in aneforge/graph.py 
717
718
719
720
721
722
723
724
def pixel_shuffle(x: Tensor, r: int) -> Tensor:
  """Depth-to-space upscale (PyTorch `nn.PixelShuffle`): `[N,C*r*r,H,W] -> [N,C,H*r,W*r]`. Fused (no cut)."""
  if len(x.shape) != 4:
    raise ValueError(f"pixel_shuffle expects 4D [N,C,H,W], got {x.shape}")
  N, C2, H, W = x.shape
  if C2 % (r * r):
    raise ValueError(f"pixel_shuffle: channels {C2} not divisible by r*r={r*r}")
  return Tensor((N, C2 // (r * r), H * r, W * r), "pixel_shuffle", [x], {"r": int(r)})

pixel_unshuffle 

pixel_unshuffle(x: Tensor, r: int) -> Tensor

Space-to-depth (PyTorch nn.PixelUnshuffle): [N,C,H*r,W*r] -> [N,C*r*r,H,W]. Fused (no cut).

Source code in aneforge/graph.py 
727
728
729
730
731
732
733
734
def pixel_unshuffle(x: Tensor, r: int) -> Tensor:
  """Space-to-depth (PyTorch `nn.PixelUnshuffle`): `[N,C,H*r,W*r] -> [N,C*r*r,H,W]`. Fused (no cut)."""
  if len(x.shape) != 4:
    raise ValueError(f"pixel_unshuffle expects 4D [N,C,H,W], got {x.shape}")
  N, C, H, W = x.shape
  if H % r or W % r:
    raise ValueError(f"pixel_unshuffle: H={H},W={W} not divisible by r={r}")
  return Tensor((N, C * r * r, H // r, W // r), "pixel_unshuffle", [x], {"r": int(r)})

space_to_channel 

space_to_channel(x: Tensor, r: int) -> Tensor

Space-to-depth on the native SpaceToChannel layer (TF channel order): [N,C,H*r,W*r] -> [N,C*r*r,H,W]. Graph cut.

Source code in aneforge/graph.py 
737
738
739
740
741
742
743
744
745
746
def space_to_channel(x: Tensor, r: int) -> Tensor:
  """Space-to-depth on the native `SpaceToChannel` layer (TF channel order): `[N,C,H*r,W*r] -> [N,C*r*r,H,W]`. Graph cut."""
  if len(x.shape) != 4:
    raise ValueError(f"space_to_channel expects 4D [N,C,H,W], got {x.shape}")
  N, C, H, W = x.shape
  if N != 1:
    raise ValueError(f"space_to_channel: the native layer supports batch N=1 only; got N={N}")
  if H % r or W % r:
    raise ValueError(f"space_to_channel: H={H},W={W} not divisible by r={r}")
  return Tensor((N, C * r * r, H // r, W // r), "space_to_channel", [x], {"r": int(r)})

channel_to_space 

channel_to_space(x: Tensor, r: int) -> Tensor

Depth-to-space on the native ChannelToSpace layer (TF channel order): [N,C*r*r,H,W] -> [N,C,H*r,W*r]. Graph cut.

Source code in aneforge/graph.py 
749
750
751
752
753
754
755
756
757
758
def channel_to_space(x: Tensor, r: int) -> Tensor:
  """Depth-to-space on the native `ChannelToSpace` layer (TF channel order): `[N,C*r*r,H,W] -> [N,C,H*r,W*r]`. Graph cut."""
  if len(x.shape) != 4:
    raise ValueError(f"channel_to_space expects 4D [N,C,H,W], got {x.shape}")
  N, C2, H, W = x.shape
  if N != 1:
    raise ValueError(f"channel_to_space: the native layer supports batch N=1 only; got N={N}")
  if C2 % (r * r):
    raise ValueError(f"channel_to_space: channels {C2} not divisible by r*r={r * r}")
  return Tensor((N, C2 // (r * r), H * r, W * r), "channel_to_space", [x], {"r": int(r)})

space_to_batch 

space_to_batch(x: Tensor, bh: int, bw: int) -> Tensor

Move spatial blocks into batch (native SpaceToBatch): [N,C,H,W] -> [N*bh*bw,C,H/bh,W/bw]. Graph cut (batch grows).

Source code in aneforge/graph.py 
761
762
763
764
765
766
767
768
769
def space_to_batch(x: Tensor, bh: int, bw: int) -> Tensor:
  """Move spatial blocks into batch (native `SpaceToBatch`): `[N,C,H,W] -> [N*bh*bw,C,H/bh,W/bw]`. Graph cut (batch grows)."""
  if len(x.shape) != 4:
    raise ValueError(f"space_to_batch expects 4D [N,C,H,W], got {x.shape}")
  N, C, H, W = x.shape
  if H % bh or W % bw:
    raise ValueError(f"space_to_batch: H={H},W={W} not divisible by (bh={bh}, bw={bw})")
  return Tensor((N * bh * bw, C, H // bh, W // bw), "space_to_batch", [x],
         {"bh": int(bh), "bw": int(bw)})

batch_to_space 

batch_to_space(x: Tensor, bh: int, bw: int) -> Tensor

Move batch blocks back into space (native BatchToSpace, inverse of space_to_batch): [N*bh*bw,C,H,W] -> [N,C,H*bh,W*bw]. Graph cut; batch must divide bh*bw.

Source code in aneforge/graph.py 
772
773
774
775
776
777
778
779
780
781
def batch_to_space(x: Tensor, bh: int, bw: int) -> Tensor:
  """Move batch blocks back into space (native `BatchToSpace`, inverse of `space_to_batch`): `[N*bh*bw,C,H,W] -> [N,C,H*bh,W*bw]`. Graph cut; batch must divide `bh*bw`."""
  if len(x.shape) != 4:
    raise ValueError(f"batch_to_space expects 4D [N,C,H,W], got {x.shape}")
  B, C, H, W = x.shape
  if B % (bh * bw):
    raise ValueError(f"batch_to_space: input batch {B} not divisible by bh*bw={bh * bw} "
            f"(arch-gated on this ANE)")
  return Tensor((B // (bh * bw), C, H * bh, W * bw), "batch_to_space", [x],
         {"bh": int(bh), "bw": int(bw)})

flatten 

flatten(x: Tensor) -> Tensor

Flatten on the native Flatten layer: collapse [C,H,W] to a 1-D vector. Graph cut.

Source code in aneforge/graph.py 
784
785
786
787
788
def flatten(x: Tensor) -> Tensor:
  """Flatten on the native `Flatten` layer: collapse [C,H,W] to a 1-D vector. Graph cut."""
  if len(x.shape) != 3:
    raise ValueError(f"flatten expects 3D [C,H,W] (the native bridge layout); got {x.shape}")
  return Tensor((int(np.prod(x.shape)),), "flatten", [x])

input_view 

input_view(x: Tensor, offset: int, size: int) -> Tensor

Contiguous view x[offset:offset+size] along Width (native InputView); x flattened to 1-D -> [size]. Graph cut.

Source code in aneforge/graph.py 
791
792
793
794
795
796
def input_view(x: Tensor, offset: int, size: int) -> Tensor:
  """Contiguous view `x[offset:offset+size]` along Width (native `InputView`); `x` flattened to 1-D -> `[size]`. Graph cut."""
  W = int(np.prod(x.shape))
  if offset < 0 or size <= 0 or offset + size > W:
    raise ValueError(f"input_view: window [{offset}:{offset + size}] out of range for W={W}")
  return Tensor((size,), "input_view", [x], {"offset": int(offset), "size": int(size)})

dynamic_slice 

dynamic_slice(x: Tensor, start: int, size: int = 2) -> Tensor

Runtime-parametric slice x[start:start+size] (native DynamicSlice). Graph cut; only verified variant is Width=4, size==2.

Source code in aneforge/graph.py 
799
800
801
802
803
804
805
806
807
def dynamic_slice(x: Tensor, start: int, size: int = 2) -> Tensor:
  """Runtime-parametric slice `x[start:start+size]` (native `DynamicSlice`). Graph cut; only verified variant is Width=4, size==2."""
  W = int(np.prod(x.shape))
  if W != 4 or size != 2:
    raise ValueError("dynamic_slice: the verified ANE variant requires a length-4 "
            f"input and size==2; got W={W}, size={size}")
  if start < 0 or start + size > W:
    raise ValueError(f"dynamic_slice: window [{start}:{start + size}] out of range for W={W}")
  return Tensor((size,), "dynamic_slice", [x], {"start": int(start), "size": int(size)})

scaled_elementwise 

scaled_elementwise(x: Tensor, z: Tensor, op: str = 'Add', scale: float = 1.0) -> Tensor

scale * (x OP z) (native ScaledElementWise). op in {Add,Mult,Min,Max}; equal-size inputs. Graph cut; Sub rejected, Mult ignores scale.

Source code in aneforge/graph.py 
810
811
812
813
814
815
816
817
818
819
820
821
822
def scaled_elementwise(x: Tensor, z: Tensor, op: str = "Add", scale: float = 1.0) -> Tensor:
  """`scale * (x OP z)` (native `ScaledElementWise`). `op` in {Add,Mult,Min,Max}; equal-size inputs. Graph cut; `Sub` rejected, `Mult` ignores `scale`."""
  ops = ("Add", "Mult", "Min", "Max")
  if op not in ops:
    raise ValueError(f"scaled_elementwise: op must be one of {ops}; got {op!r} "
            f"('Sub' is rejected by the ANE ScaledElementWise layer)")
  if op == "Mult" and float(scale) != 1.0:
    raise ValueError("scaled_elementwise: the native layer ignores `scale` for op='Mult' "
            "(would silently give x*z, not scale*(x*z)); use scale=1.0 or a separate mul")
  if int(np.prod(x.shape)) != int(np.prod(z.shape)):
    raise ValueError(f"scaled_elementwise: x and z must have equal size; got {x.shape}, {z.shape}")
  return Tensor((int(np.prod(x.shape)),), "scaled_elementwise", [x, z],
         {"op": op, "scale": float(scale)})

topk 

topk(x: Tensor, k: int, largest: bool = True) -> Tensor

Top-k per row of a 2D input [C,W] (native TopK bridge, a cut). k in {3,4} is arch-gated and rejected.

Source code in aneforge/graph.py 
825
826
827
828
829
830
831
832
833
834
def topk(x: Tensor, k: int, largest: bool = True) -> Tensor:
  """Top-`k` per row of a 2D input [C,W] (native TopK bridge, a cut). `k` in {3,4} is arch-gated and rejected."""
  if len(x.shape) != 2:
    raise ValueError(f"topk: only 2D [C,W] inputs are supported; got {x.shape}")
  C, W = x.shape
  if not (1 <= k <= W):
    raise ValueError(f"topk: k={k} out of range [1, {W}]")
  if k in (3, 4):
    raise ValueError(f"topk: k={k} is arch-gated on this ANE (ANECCompile fails for k in {{3,4}})")
  return Tensor((C, k), "topk", [x], {"k": int(k), "largest": bool(largest)})

sort 

sort(x: Tensor, descending: bool = False, return_indices: bool = False) -> Tensor

Sort each row of a 2D input [C,W] along Width (native Sort bridge, a cut). return_indices gives fp16-encoded argsort indices.

Source code in aneforge/graph.py 
837
838
839
840
841
842
def sort(x: Tensor, descending: bool = False, return_indices: bool = False) -> Tensor:
  """Sort each row of a 2D input [C,W] along Width (native Sort bridge, a cut). `return_indices` gives fp16-encoded argsort indices."""
  if len(x.shape) != 2:
    raise ValueError(f"sort: only 2D [C,W] inputs are supported; got {x.shape}")
  return Tensor(x.shape, "sort", [x],
         {"descending": bool(descending), "return_indices": bool(return_indices)})

cross_product 

cross_product(a: Tensor, b: Tensor) -> Tensor

3-vector cross product cross(a,b) (native CrossProduct layer). Both inputs length-3; returns (3,). Graph cut.

Source code in aneforge/graph.py 
845
846
847
848
849
def cross_product(a: Tensor, b: Tensor) -> Tensor:
  """3-vector cross product `cross(a,b)` (native CrossProduct layer). Both inputs length-3; returns (3,). Graph cut."""
  if int(np.prod(a.shape)) != 3 or int(np.prod(b.shape)) != 3:
    raise ValueError(f"cross_product: both inputs must have 3 elements; got {a.shape}, {b.shape}")
  return Tensor((3,), "cross_product", [a, b])

cross_correlation 

cross_correlation(x: Tensor, template: Tensor) -> Tensor

Valid (no-flip) cross-correlation of map x [H,W] with template [Th,Tw] (native CrossCorrelation): y[i,j] = sum x[i+u,j+v]*template[u,v] -> [H-Th+1, W-Tw+1]. Graph cut.

Source code in aneforge/graph.py 
852
853
854
855
856
857
858
859
860
def cross_correlation(x: Tensor, template: Tensor) -> Tensor:
  """Valid (no-flip) cross-correlation of map `x` [H,W] with `template` [Th,Tw] (native CrossCorrelation): `y[i,j] = sum x[i+u,j+v]*template[u,v]` -> [H-Th+1, W-Tw+1]. Graph cut."""
  if len(x.shape) != 2 or len(template.shape) != 2:
    raise ValueError(f"cross_correlation: x and template must be 2D; got {x.shape}, {template.shape}")
  H, W = x.shape
  Th, Tw = template.shape
  if Th > H or Tw > W:
    raise ValueError(f"cross_correlation: template {template.shape} larger than map {x.shape}")
  return Tensor((H - Th + 1, W - Tw + 1), "cross_correlation", [x, template])

cost_volume 

cost_volume(aux: Tensor, ref: Tensor, disparity_range: int = 1) -> Tensor

L1 stereo/flow matching cost (native CostVolume). aux length-Wa, ref length-Wr (Wr>=Wa+R) -> (R+1,Wa) with cost[d,x]=|aux[x]-ref[x+d]|. Graph cut.

Source code in aneforge/graph.py 
863
864
865
866
867
868
869
870
871
def cost_volume(aux: Tensor, ref: Tensor, disparity_range: int = 1) -> Tensor:
  """L1 stereo/flow matching cost (native CostVolume). `aux` length-Wa, `ref` length-Wr (Wr>=Wa+R) -> `(R+1,Wa)` with `cost[d,x]=|aux[x]-ref[x+d]|`. Graph cut."""
  Wa, Wr = int(np.prod(aux.shape)), int(np.prod(ref.shape))
  R = int(disparity_range)
  if R < 0:
    raise ValueError(f"cost_volume: disparity_range must be >= 0; got {R}")
  if Wr < Wa + R:
    raise ValueError(f"cost_volume: ref width {Wr} must be >= aux width {Wa} + disparity_range {R}")
  return Tensor((R + 1, Wa), "cost_volume", [aux, ref], {"disparity_range": R})

fps 

fps(points: Tensor, k: int) -> Tensor

Furthest-point sampling: greedily pick k far-apart points (native FurthestPointSampling, L2 only). points [N,3] -> [k,3] centroids. Graph cut.

Source code in aneforge/graph.py 
874
875
876
877
878
879
880
881
882
883
def fps(points: Tensor, k: int) -> Tensor:
  """Furthest-point sampling: greedily pick `k` far-apart points (native FurthestPointSampling, L2 only). `points` [N,3] -> [k,3] centroids. Graph cut."""
  if len(points.shape) != 2 or points.shape[1] != 3:
    raise ValueError(f"fps: points must be [N, 3]; got {points.shape}")
  N = points.shape[0]
  if not (1 <= k <= N):
    raise ValueError(f"fps: k={k} out of range [1, {N}]")
  if k > 1024 or N > 8192:
    raise ValueError(f"fps: arch limits are k<=1024, N<=8192; got k={k}, N={N}")
  return Tensor((k, 3), "fps", [points], {"k": int(k)})

radius_search(points: Tensor, centroids: Tensor, radius: float) -> Tensor

L2 ball-query membership (native RadiusSearch): 1 iff point within radius of centroid. points [N,3], centroids [Nc,3] -> [N,Nc] 0/1. Graph cut.

Source code in aneforge/graph.py 
886
887
888
889
890
891
892
893
def radius_search(points: Tensor, centroids: Tensor, radius: float) -> Tensor:
  """L2 ball-query membership (native RadiusSearch): 1 iff point within `radius` of centroid. `points` [N,3], `centroids` [Nc,3] -> [N,Nc] 0/1. Graph cut."""
  if len(points.shape) != 2 or points.shape[1] != 3:
    raise ValueError(f"radius_search: points must be [N, 3]; got {points.shape}")
  if len(centroids.shape) != 2 or centroids.shape[1] != 3:
    raise ValueError(f"radius_search: centroids must be [Nc, 3]; got {centroids.shape}")
  N, Nc = points.shape[0], centroids.shape[0]
  return Tensor((N, Nc), "radius_search", [points, centroids], {"radius": float(radius)})

minmax_norm 

minmax_norm(x: Tensor, dimension: str = 'Width', eps: float = 0.0001) -> Tensor

Min-max normalize (x-min)/(max-min+eps) over dimension (native MinMaxNormalization). x [1,C,H,W]; "Width"/"Height" only. Graph cut.

Source code in aneforge/graph.py 
896
897
898
899
900
901
902
903
def minmax_norm(x: Tensor, dimension: str = "Width", eps: float = 1e-4) -> Tensor:
  """Min-max normalize `(x-min)/(max-min+eps)` over `dimension` (native MinMaxNormalization). `x` [1,C,H,W]; "Width"/"Height" only. Graph cut."""
  if len(x.shape) != 4 or x.shape[0] != 1:
    raise ValueError(f"minmax_norm: expects [1,C,H,W]; got {x.shape}")
  if dimension not in ("Width", "Height"):
    raise ValueError(f"minmax_norm: dimension must be 'Width' or 'Height' "
            f"('Channel' is arch-gated on this ANE); got {dimension!r}")
  return Tensor(x.shape, "minmax_norm", [x], {"dimension": dimension, "eps": float(eps)})

lrn 

lrn(x: Tensor, alpha: float = 1.0, beta: float = 0.75, k: float = 1.0) -> Tensor

Cross-channel LRN (AlexNet) on the native LocalResponseNormalization layer (Channel mode). x [1,C,H,W]; graph cut. Window is a clipped local channel window of size N=C. Arch-gated: C<=15 only.

Source code in aneforge/graph.py 
906
907
908
909
910
911
912
913
914
def lrn(x: Tensor, alpha: float = 1.0, beta: float = 0.75, k: float = 1.0) -> Tensor:
  """Cross-channel LRN (AlexNet) on the native LocalResponseNormalization layer (Channel mode). `x` [1,C,H,W]; graph cut. Window is a clipped local channel window of size N=C. Arch-gated: C<=15 only."""
  if len(x.shape) != 4 or x.shape[0] != 1:
    raise ValueError(f"lrn: expects [1,C,H,W]; got {x.shape}")
  C = x.shape[1]
  if C > 15:
    raise ValueError(f"lrn: C={C} is arch-gated on this ANE (LocalResponseNormalization "
            f"with KernelChannel=C fails ANECCompile for C>=16); got C={C}")
  return Tensor(x.shape, "lrn", [x], {"alpha": float(alpha), "beta": float(beta), "k": float(k)})

mha 

mha(x: Tensor, Wq, bq, Wk, bk, Wv, bv, Wo, bo, n_heads: int, mask=None) -> Tensor

Multi-head self-attention on x [S,D]. Weights [out,in]; biases [D] or None. mask, when given, is an additive score bias broadcast to [H,S,S] and sliced along the query axis -- pass [1,S,S] (or [S,S]) with -inf/-1e4 at padded key columns for a key-padding mask (lets a padded batch share one program without pad tokens corrupting the real ones).

Source code in aneforge/graph.py 
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
def mha(x: Tensor, Wq, bq, Wk, bk, Wv, bv, Wo, bo, n_heads: int, mask=None) -> Tensor:
  """Multi-head self-attention on `x` [S,D]. Weights [out,in]; biases [D] or None. `mask`, when given,
  is an additive score bias broadcast to [H,S,S] and sliced along the query axis -- pass `[1,S,S]` (or
  `[S,S]`) with -inf/-1e4 at padded key columns for a key-padding mask (lets a padded batch share one
  program without pad tokens corrupting the real ones)."""
  S, D = x.shape
  if D % n_heads:
    raise ValueError(f"mha: D={D} not divisible by n_heads={n_heads}")
  dh = D // n_heads
  q, k, v = x.linear(Wq, bq), x.linear(Wk, bk), x.linear(Wv, bv)
  qh, kh, vh = _heads(q, n_heads, dh), _heads(k, n_heads, dh), _heads(v, n_heads, dh)
  kt = kh.transpose([0, 2, 1])
  scale = 1.0 / dh ** 0.5
  # query-tiling: [tile, S] score tiles per head instead of the full [H, S, S] matrix
  from . import _optimize as _opt
  n_tiles = _opt.attention_tiles(S, n_heads, dh)
  o = _tiled_attention(qh, kt, vh, scale, n_tiles, seq_axis=1, mask=mask)  # [H, S, dh]
  o = o.transpose([1, 0, 2]).reshape(S, D)
  return o.linear(Wo, bo)

cross_attention 

cross_attention(x: Tensor, context: Tensor, Wq, Wk, Wv, Wo, n_heads: int, bq=None, bk=None, bv=None, bo=None) -> Tensor

Cross-attention: queries from x [S,D], keys/values from context [T,Dctx]. Wq:[D,D]; Wk,Wv:[D,Dctx]; Wo:[D,D].

Source code in aneforge/graph.py 
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
def cross_attention(x: Tensor, context: Tensor, Wq, Wk, Wv, Wo, n_heads: int,
          bq=None, bk=None, bv=None, bo=None) -> Tensor:
  """Cross-attention: queries from `x` [S,D], keys/values from `context` [T,Dctx]. Wq:[D,D]; Wk,Wv:[D,Dctx]; Wo:[D,D]."""
  S, D = x.shape
  T = context.shape[0]
  dh = D // n_heads
  qh = _heads(x.linear(Wq, bq), n_heads, dh)                              # [H,S,dh]
  kh = _heads(context.linear(Wk, bk), n_heads, dh)                        # [H,T,dh]
  vh = _heads(context.linear(Wv, bv), n_heads, dh)
  kt = kh.transpose([0, 2, 1])                                            # [H,dh,T]
  scale = 1.0 / dh ** 0.5
  # query-tiling when both query and context are long (small-T stays single-shot)
  from . import _optimize as _opt
  n_tiles = _opt.attention_tiles(S, n_heads, dh, T=T) if (S >= 768 and T >= 512) else 1
  o = _tiled_attention(qh, kt, vh, scale, n_tiles, seq_axis=1)  # [H, S, dh]
  o = o.transpose([1, 0, 2]).reshape(S, D)
  return o.linear(Wo, bo)

sdpa 

sdpa(q: Tensor, k: Tensor, v: Tensor, scale: float | None = None, is_causal: bool = False, attn_mask: 'Tensor | None' = None) -> Tensor

Scaled-dot-product attention via the native fused-attention layer (ANECSDPALayerDesc) inside the reliable regime, else the fused decomposition. q/k/v: [1,heads,seq,d_head] fp16; native use is a graph cut. is_causal=True is native (causal mask on the 5th bottom).

Source code in aneforge/graph.py 
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
def sdpa(q: Tensor, k: Tensor, v: Tensor, scale: float | None = None,
    is_causal: bool = False, attn_mask: "Tensor | None" = None) -> Tensor:
  """Scaled-dot-product attention via the native fused-attention layer (ANECSDPALayerDesc) inside the reliable regime, else the fused decomposition. q/k/v: [1,heads,seq,d_head] fp16; native use is a graph cut. `is_causal=True` is native (causal mask on the 5th bottom)."""
  # K,V share shape (cached seq); Q's seq may differ (KV-cache decode). Q,K share H+D.
  if not (len(q.shape) == 4 == len(k.shape) == len(v.shape)):
    raise ValueError(f"af.sdpa expects 4D q,k,v of [1,H,S,D]; got {q.shape}, {k.shape}, {v.shape}")
  if k.shape != v.shape:
    raise ValueError(f"af.sdpa: k,v must share shape (the cached sequence); got {k.shape}, {v.shape}")
  if q.shape[1] != k.shape[1] or q.shape[3] != k.shape[3]:
    raise ValueError(f"af.sdpa: q,k must share H (heads) and D (embedding); got {q.shape}, {k.shape}")
  if q.shape[0] != 1 or k.shape[0] != 1:
    raise ValueError(f"af.sdpa: batch must be 1 (native layer); got {q.shape[0]}/{k.shape[0]}")
  if is_causal and q.shape[2] != k.shape[2]:
    raise ValueError("af.sdpa: is_causal requires equal q/k seq (prefill); for KV-cache "
            "decode (seq_q < seq_kv) the new tokens attend to all cached k/v - "
            "pass is_causal=False (or a runtime attn_mask).")
  if attn_mask is not None:
    if is_causal:
      raise ValueError("af.sdpa: pass either is_causal or an explicit attn_mask, not both")
    # attn_mask is a runtime additive bias on the 5th bottom: ONE shared plane [1,1,Sq,Skv].
    # Per-head and query-broadcast forms are mis-applied -- reject both.
    if (len(attn_mask.shape) != 4 or attn_mask.shape[0] != 1 or attn_mask.shape[1] != 1
        or attn_mask.shape[2] != q.shape[2] or attn_mask.shape[3] != k.shape[2]):
      raise ValueError(
        f"af.sdpa: attn_mask must be a single shared plane [1,1,Sq,Skv]="
        f"{[1, 1, q.shape[2], k.shape[2]]} (one mask for all heads, full query axis); "
        f"got {list(attn_mask.shape)}. Per-head masks (H>1) and query-broadcast "
        f"(Sq-axis=1 while q_seq>1) are not supported by the native layer.")
  _scale: float = scale if scale is not None else 1.0 / q.shape[-1] ** 0.5
  seq = max(q.shape[2], k.shape[2])               # attention spans the K/V (cached) length
  both = min(q.shape[2], k.shape[2])              # native layer breaks when BOTH axes are large
  native_ok = both < SDPA_NATIVE_MIN_BOTH and seq <= SDPA_NATIVE_MAX_SEQ
  if not native_ok:
    if is_causal:
      # decomposition has no causal mask and native is unreliable here -- refuse
      raise NotImplementedError(
        f"af.sdpa: causal attention at min(q,k)seq={both} (>= {SDPA_NATIVE_MIN_BOTH}) "
        f"or seq={seq} (> {SDPA_NATIVE_MAX_SEQ}) is outside the reliable native regime "
        f"and the causal decomposition is not wired; chunk the query so each tile's "
        f"min(seq) < {SDPA_NATIVE_MIN_BOTH}.")
    # non-causal: the fused decomposition (query-tiled for a large query axis)
    kt = k.transpose([0, 1, 3, 2])
    from . import _optimize as _opt
    n_tiles = _opt.attention_tiles(q.shape[2], q.shape[1], q.shape[3], T=k.shape[2])
    return _tiled_attention(q, kt, v, float(_scale), n_tiles, seq_axis=2, mask=attn_mask)
  if attn_mask is not None:                       # runtime mask rides the 5th bottom (stays native)
    return Tensor(q.shape, "sdpa", [q, k, v, attn_mask], {"scale": float(_scale), "masked": True})
  return Tensor(q.shape, "sdpa", [q, k, v], {"scale": float(_scale), "causal": bool(is_causal)})

geglu 

geglu(x: Tensor, W, b) -> Tensor

GEGLU FFN gate: split the [2*Dff,D] projection into value/gate halves; out = value * gelu(gate).

Source code in aneforge/graph.py 
1043
1044
1045
1046
1047
1048
1049
def geglu(x: Tensor, W, b) -> Tensor:
  """GEGLU FFN gate: split the [2*Dff,D] projection into value/gate halves; out = value * gelu(gate)."""
  W = np.asarray(W); Dff = W.shape[0] // 2
  bv = bg = None
  if b is not None:
    b = np.asarray(b); bv, bg = b[:Dff], b[Dff:]
  return x.linear(W[:Dff], bv) * x.linear(W[Dff:], bg).gelu()