Skip to content

Applied math

Self-contained numerical toolkits built over the public operators: iterative linear algebra, an FFT carried as real/imaginary pairs, special functions, a general einsum, and signal processing.

Linear algebra

linalg

aneforge.linalg - iterative solvers and small-n direct factorizations on the ANE.

conjugate_gradient 

conjugate_gradient(A, b, iters: int = 20, x0=None, refine: int = 0)

Solve A x = b for SPD A by CG with fixed iters (recurrence unrolls into one program); refine adds residual-correction rounds.

Source code in aneforge/linalg.py 
 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
def conjugate_gradient(A, b, iters: int = 20, x0=None, refine: int = 0):
  """Solve A x = b for SPD A by CG with fixed `iters` (recurrence unrolls into one program); `refine` adds residual-correction rounds."""
  A0 = np.asarray(A, np.float64); b0 = np.asarray(b, np.float64).reshape(-1)
  n = A0.shape[0]
  # symmetric Jacobi preconditioner: A~ = Dh^{-1} A Dh^{-1}, b~ = Dh^{-1} b,
  # x = Dh^{-1} y;  Dh = sqrt(diag(A)).  A~ symmetric, so A~ p (row) is p @ A~.
  dh = np.sqrt(np.abs(A0.diagonal())); dh[dh == 0] = 1.0
  A16 = f16((A0 / dh[:, None]) / dh[None, :])
  b16 = f16(b0 / dh).reshape(1, n)
  ones = np.ones((n, 1), f16)

  bT = af.input((1, n))
  x0T = af.input((1, n)) if x0 is not None else None
  dot = lambda u, v: (u * v) @ ones                       # ANE matmul dot (saturates above ~32752 = fp16_max/2; scaling keeps sums in range)

  def cg_block(x, r, K):
    """K unrolled CG steps from residual r (p=r), accumulating into x."""
    p = r
    rs = dot(r, r)
    for _ in range(K):
      Ap = p @ A16                                     # ANE GEMV  A~ p
      alpha = rs / dot(p, Ap)                          # [1,1] / [1,1]
      x = x + alpha * p                                # broadcast axpy
      r = r - alpha * Ap
      rs_new = dot(r, r)
      p = r + (rs_new / rs) * p                        # beta = rs_new/rs
      rs = rs_new
    return x, r

  if x0T is None:
    x, r = bT * 0.0, bT                                  # x=0 -> r0 = b
  else:
    x, r = x0T, bT - x0T @ A16                           # r0 = b - A~ x0
  x, r = cg_block(x, r, iters)
  for _ in range(refine):
    r = bT - x @ A16                                     # residual correction
    dx, _ = cg_block(bT * 0.0, r, max(4, iters // 2))
    x = x + dx

  feeds = [b16] if x0T is None else [b16, (np.asarray(x0, np.float64).reshape(-1) * dh).astype(f16).reshape(1, n)]
  y = _solve_once(x, *feeds).ravel()
  if not np.isfinite(y).all(): y = np.zeros(n)   # high cond: iterates can overflow fp16
  return (y / dh).astype(np.float32)

jacobi 

jacobi(A, b, iters: int = 60, x0=None)

Jacobi iteration x += D^{-1}(b - A x), fixed iters; converges for diagonally-dominant A.

Source code in aneforge/linalg.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def jacobi(A, b, iters: int = 60, x0=None):
  """Jacobi iteration x += D^{-1}(b - A x), fixed `iters`; converges for diagonally-dominant A."""
  A16 = np.asarray(A, f16); b0 = np.asarray(b, np.float64).reshape(-1)
  n = A16.shape[0]
  AT16 = np.ascontiguousarray(A16.T)                     # x @ A^T = (A x) as a row
  d = np.asarray(A16, np.float64).diagonal().copy(); d[d == 0] = 1.0
  dinv = (1.0 / d).astype(f16).reshape(1, n)

  bT = af.input((1, n)); dinvT = af.input((1, n))
  x0T = af.input((1, n)) if x0 is not None else None
  x = bT * 0.0 if x0T is None else x0T
  for _ in range(iters):
    x = x + (bT - x @ AT16) * dinvT                    # x + D^{-1}(b - A x)
  feeds = [b0.astype(f16).reshape(1, n), dinv]
  if x0T is not None: feeds.append(np.asarray(x0, f16).reshape(1, n))
  return _solve_once(x, *feeds).ravel().astype(np.float32)

gauss_seidel 

gauss_seidel(A, b, iters: int = 60, x0=None)

Gauss-Seidel iteration, fixed iters; serial sweep is host-side, residual GEMV on the ANE.

Source code in aneforge/linalg.py 
135
136
137
138
139
140
141
142
143
144
145
146
def gauss_seidel(A, b, iters: int = 60, x0=None):
  """Gauss-Seidel iteration, fixed `iters`; serial sweep is host-side, residual GEMV on the ANE."""
  A16 = np.asarray(A, f16); b16 = np.asarray(b, f16).reshape(-1)
  Af = np.asarray(A16, np.float64); bf = b16.astype(np.float64)
  n = A16.shape[0]
  x = np.zeros(n, np.float64) if x0 is None else np.asarray(x0, np.float64).reshape(-1).copy()
  L = np.tril(Af, -1); U = np.triu(Af, 1); d = Af.diagonal()
  for _ in range(iters):
    for i in range(n):
      s = L[i, :i] @ x[:i] + U[i, i + 1:] @ x[i + 1:]
      x[i] = (bf[i] - s) / d[i]
  return x.astype(np.float32)

iterative_refine 

iterative_refine(A, b, x0, iters: int = 3, inner: int = 60)

Sharpen a solve x0 by fixed-K residual correction (r = b - A x, solve A dx = r, x += dx); inner is a Richardson sweep.

Source code in aneforge/linalg.py 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def iterative_refine(A, b, x0, iters: int = 3, inner: int = 60):
  """Sharpen a solve x0 by fixed-K residual correction (r = b - A x, solve A dx = r, x += dx); inner is a Richardson sweep."""
  A0 = np.asarray(A, np.float64); b0 = np.asarray(b, np.float64).reshape(-1)
  n = A0.shape[0]
  # same symmetric Jacobi scaling as CG, to keep the fp16 GEMV products in range
  dh = np.sqrt(np.abs(A0.diagonal())); dh[dh == 0] = 1.0
  A16 = f16((A0 / dh[:, None]) / dh[None, :])
  b16 = f16(b0 / dh).reshape(1, n)
  omega = _spectral_omega(A16)

  bT = af.input((1, n)); x0T = af.input((1, n))
  x = x0T
  for _ in range(iters):
    r = bT - x @ A16                                    # residual (ANE GEMV + sub)
    dx = bT * 0.0
    for _ in range(inner):
      dx = dx + (r - dx @ A16) * omega                # Richardson, omega is a scalar
    x = x + dx
  x0s = (np.asarray(x0, np.float64).reshape(-1) * dh).astype(f16).reshape(1, n)
  y = _solve_once(x, b16, x0s).ravel()
  return (y / dh).astype(np.float32)

randomized_svd 

randomized_svd(A, k: int, oversample: int = 5, power_iters: int = 2)

Truncated SVD of A [m,n] via Halko-Martinsson-Tropp range finding; returns (U[:, :k], S[:k], Vt[:k, :]).

Source code in aneforge/linalg.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def randomized_svd(A, k: int, oversample: int = 5, power_iters: int = 2):
  """Truncated SVD of A [m,n] via Halko-Martinsson-Tropp range finding; returns (U[:, :k], S[:k], Vt[:k, :])."""
  A16 = np.asarray(A, f16)
  m, n = A16.shape
  l = min(k + oversample, n)
  rng = np.random.default_rng(0)
  Omega = rng.standard_normal((n, l)).astype(f16)                        # HOST RNG (a constant)

  Y = _ane_gemm(A16, Omega)                                              # ANE: A @ Omega -> [m,l]
  Q16 = np.linalg.qr(np.asarray(Y, np.float64))[0].astype(f16)           # HOST QR -> O(1) basis
  for _ in range(power_iters):
    AtQ = _ane_gemm(A16, Q16, transpose_a=True)                        # ANE: A^T @ Q -> [n,l]
    Y = _ane_gemm(A16, AtQ)                                            # ANE: A   @ AtQ -> [m,l]
    Q16 = np.linalg.qr(np.asarray(Y, np.float64))[0].astype(f16)       # HOST re-orthonormalize

  B = _ane_gemm(A16, Q16, transpose_a=True).T                           # ANE: (A^T @ Q)^T = Q^T @ A -> [l,n]
  Ub, S, Vt = np.linalg.svd(np.asarray(B, np.float64), full_matrices=False)  # HOST small SVD
  U = _ane_gemm(Q16, Ub.astype(f16))                                     # ANE: Q @ Ub -> [m,l]
  return (np.asarray(U, np.float32)[:, :k],
          S[:k].astype(np.float32),
          np.asarray(Vt, np.float32)[:k, :])

pca 

pca(X, k: int, oversample: int = 5, power_iters: int = 2)

Principal components of X [samples, features] via randomized_svd of the centered data; returns (components, singular_values, mean).

Source code in aneforge/linalg.py 
199
200
201
202
203
204
205
def pca(X, k: int, oversample: int = 5, power_iters: int = 2):
  """Principal components of X [samples, features] via randomized_svd of the centered data; returns (components, singular_values, mean)."""
  Xf = np.asarray(X, np.float64)
  mean = Xf.mean(0)
  Xc = (Xf - mean).astype(f16)                                           # HOST centering
  _, S, Vt = randomized_svd(Xc, k, oversample, power_iters)              # ANE matmuls inside
  return Vt[:k].astype(np.float32), S[:k].astype(np.float32), mean.astype(np.float32)

least_squares 

least_squares(A, b, iters: int = 40, refine: int = 2)

Solve min_x ||A x - b||_2 via the normal equations A^T A x = A^T b by CG + refinement (forming A^T A squares cond(A)).

Source code in aneforge/linalg.py 
210
211
212
213
214
215
216
217
218
219
220
221
def least_squares(A, b, iters: int = 40, refine: int = 2):
  """Solve min_x ||A x - b||_2 via the normal equations A^T A x = A^T b by CG + refinement (forming A^T A squares cond(A))."""
  A16 = np.asarray(A, f16); b16 = np.asarray(b, f16).reshape(-1)
  m, n = A16.shape
  AtA = _ane_gemm(A16, A16, transpose_a=True)                            # ANE: A^T A -> [n,n]
  # A^T b  via matmul: (A^T) @ b  ==  ( [n,m] @ [m,1] )
  At = af.input((m, n))
  bt = af.input((m, 1))
  net = af.compile(At.transpose([1, 0]) @ bt)                            # ANE GEMV
  Atb = net(A16, b16.reshape(m, 1).astype(f16)).reshape(n)
  net.release()
  return conjugate_gradient(AtA, Atb, iters=iters, refine=refine)

lsqr 

lsqr(A, b, iters: int = 80)

Solve min_x ||A x - b||_2 by LSQR (Golub-Kahan bidiagonalization), fully on the ANE; square systems only to cond~1e1 (use gmres beyond).

Source code in aneforge/linalg.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
def lsqr(A, b, iters: int = 80):
  """Solve min_x ||A x - b||_2 by LSQR (Golub-Kahan bidiagonalization), fully on the ANE; square systems only to cond~1e1 (use gmres beyond)."""
  A0 = np.asarray(A, f16); m, n = A0.shape; AT = np.ascontiguousarray(A0.T)
  onem = np.ones((m, 1), f16); onen = np.ones((n, 1), f16)
  bT = af.input((1, m))
  nrm = lambda z, o: ((z * z) @ o).sqrt()
  Ax = lambda v: v @ AT                                # [1,n]@[n,m] = A@v
  ATx = lambda u: u @ A0                               # [1,m]@[m,n] = A^T@u
  beta = nrm(bT, onem); u = bT / beta
  v = ATx(u); alpha = nrm(v, onen); v = v / alpha
  w = v; x = v * 0.0; phib = beta; rhob = alpha
  for _ in range(iters):
    u2 = Ax(v) - alpha * u; beta = nrm(u2, onem); u = u2 / beta
    v2 = ATx(u) - beta * v; alpha = nrm(v2, onen); v = v2 / alpha
    rho = (rhob * rhob + beta * beta).sqrt(); c = rhob / rho; s = beta / rho
    theta = s * alpha; rhob = (c * alpha) * -1.0; phi = c * phib; phib = s * phib
    x = x + (phi / rho) * w; w = v - (theta / rho) * w
  return _solve_once(x, np.asarray(b, f16).reshape(1, m)).ravel().astype(np.float32)

dominant_eig 

dominant_eig(A, iters: int = 60, seed: int = 0)

Dominant eigenpair of symmetric A by power iteration, fully on the ANE (in-graph Rayleigh quotient); returns (lambda, eigenvector).

Source code in aneforge/linalg.py 
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
def dominant_eig(A, iters: int = 60, seed: int = 0):
  """Dominant eigenpair of symmetric A by power iteration, fully on the ANE (in-graph Rayleigh quotient); returns (lambda, eigenvector)."""
  A0 = np.asarray(A, f16); n = A0.shape[0]; AT = np.ascontiguousarray(A0.T)
  onen = np.ones((n, 1), f16)
  x0 = af.input((1, n))
  nrm = lambda z: ((z * z) @ onen).sqrt()
  Ax = lambda v: v @ AT
  x = x0 / nrm(x0)
  for _ in range(iters):
    y = Ax(x); x = y / nrm(y)
  Axf = Ax(x)
  lam = ((x * Axf) @ onen) / ((x * x) @ onen)          # Rayleigh quotient [1,1]
  out = af.concat([lam, x], axis=1)                    # [1, 1+n]
  seedv = np.random.default_rng(seed).standard_normal((1, n)).astype(f16)
  r = _solve_once(out, seedv).ravel()
  return float(r[0]), r[1:].astype(np.float32)

gmres 

gmres(A, b)

Solve A x = b for general A by one full GMRES(n) cycle, fully on the ANE; more accurate than lsqr on square systems but heavier.

Source code in aneforge/linalg.py 
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
def gmres(A, b):
  """Solve A x = b for general A by one full GMRES(n) cycle, fully on the ANE; more accurate than lsqr on square systems but heavier."""
  A0 = np.asarray(A, f16); n = A0.shape[0]; m = n
  AT = np.ascontiguousarray(A0.T); onen = np.ones((n, 1), f16)
  bT = af.input((1, n))
  nrm = lambda z: ((z * z) @ onen).sqrt()
  Ax = lambda v: v @ AT
  dot = lambda a, c: (a * c) @ onen
  Q = [bT / nrm(bT)]; H = {}; g = [nrm(bT)] + [None] * m; cs = [None] * m; sn = [None] * m
  for j in range(m):
    w = Ax(Q[j])
    for _p in range(2):                                  # double MGS, all in-graph
      for i in range(j + 1):
        h = dot(Q[i], w); H[(i, j)] = h if (i, j) not in H else H[(i, j)] + h; w = w - h * Q[i]
    hn = nrm(w); Hjp = hn
    for i in range(j):                                   # apply prior Givens to column j
      t = cs[i] * H[(i, j)] + sn[i] * H[(i + 1, j)]
      H[(i + 1, j)] = sn[i] * (H[(i, j)] * -1.0) + cs[i] * H[(i + 1, j)]; H[(i, j)] = t
    H[(j + 1, j)] = Hjp
    d = (H[(j, j)] * H[(j, j)] + Hjp * Hjp).sqrt(); cs[j] = H[(j, j)] / d; sn[j] = Hjp / d
    H[(j, j)] = cs[j] * H[(j, j)] + sn[j] * Hjp
    g[j + 1] = sn[j] * (g[j] * -1.0); g[j] = cs[j] * g[j]
    Q.append(w / hn)
  y = [None] * m                                           # back-substitution R y = g
  for i in range(m - 1, -1, -1):
    acc = g[i]
    for jj in range(i + 1, m):
      acc = acc - H[(i, jj)] * y[jj]
    y[i] = acc / H[(i, i)]
  x = Q[0] * 0.0
  for i in range(m): x = x + y[i] * Q[i]
  return _solve_once(x, np.asarray(b, f16).reshape(1, n)).ravel().astype(np.float32)

dominant_svd 

dominant_svd(A, iters: int = 60, seed: int = 0)

Dominant singular triple by power iteration on A^T A, fully on the ANE (half-step renorm holds fp16 range); returns (sigma1, u1, v1).

Source code in aneforge/linalg.py 
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
def dominant_svd(A, iters: int = 60, seed: int = 0):
  """Dominant singular triple by power iteration on A^T A, fully on the ANE (half-step renorm holds fp16 range); returns (sigma1, u1, v1)."""
  A0 = np.asarray(A, f16); m, n = A0.shape
  AT = np.ascontiguousarray(A0.T); onen = np.ones((n, 1), f16); onem = np.ones((m, 1), f16)
  v0 = af.input((1, n))
  nn = lambda z: ((z * z) @ onen).sqrt()
  nm = lambda z: ((z * z) @ onem).sqrt()
  Ax = lambda v: v @ AT                                  # [1,n]@[n,m] = A v
  ATx = lambda u: u @ A0                                 # [1,m]@[m,n] = A^T u
  v = v0 / nn(v0)
  for _ in range(iters):
    Av = Ax(v); Av = Av / nm(Av)                          # half-step renorm: holds O(1)
    v = ATx(Av); v = v / nn(v)
  Av = Ax(v); sig = nm(Av)                               # sigma1 = ||A v1||  [1,1]
  u = Av / sig
  out = af.concat([sig, u, v], axis=1)                   # [1, 1+m+n]
  seedv = np.random.default_rng(seed).standard_normal((1, n)).astype(f16)
  r = _solve_once(out, seedv).ravel()
  return float(r[0]), r[1:1 + m].astype(np.float32), r[1 + m:].astype(np.float32)

eigh 

eigh(A, sweeps: int = 8, iterate: bool = False)

All eigenvalues of symmetric A by fixed-sweep cyclic Jacobi, ascending; iterate=True host-loops one compiled sweep (reaches larger n).

Source code in aneforge/linalg.py 
356
357
358
359
360
361
362
363
364
365
366
367
368
def eigh(A, sweeps: int = 8, iterate: bool = False):
  """All eigenvalues of symmetric A by fixed-sweep cyclic Jacobi, ascending; `iterate=True` host-loops one compiled sweep (reaches larger n)."""
  A0 = np.asarray(A, f16); n = A0.shape[0]
  if iterate:
    net = af.compile(_jacobi_sweep(af.input((n, n)), n), _check_precision=False)
    M = A0
    for _ in range(sweeps): M = net(M).astype(f16)
    net.release()
    return np.sort(np.diag(M.astype(np.float64))).astype(np.float32)
  M = af.input((n, n))
  for _ in range(sweeps): M = _jacobi_sweep(M, n)
  out = _solve_once(M, A0)                                                 # final ~diagonal matrix
  return np.sort(np.diag(out)).astype(np.float32)

svd 

svd(A, sweeps: int = 8)

All singular values of A (descending), fully on the ANE: Gram matrix via one GEMM, then its spectrum via cyclic-Jacobi eigh.

Source code in aneforge/linalg.py 
371
372
373
374
375
376
377
378
def svd(A, sweeps: int = 8):
  """All singular values of A (descending), fully on the ANE: Gram matrix via one GEMM, then its spectrum via cyclic-Jacobi eigh."""
  A16 = np.asarray(A, f16); m, n = A16.shape
  # form the SMALLER Gram matrix (eigh's graph is O(dim^3)): same nonzero eigenvalues.
  M = A16 if m >= n else np.ascontiguousarray(A16.T)
  G = _ane_gemm(M, M, transpose_a=True)                  # ANE: M^T M  [min(m,n), min(m,n)]
  ev = eigh(G, sweeps=sweeps)
  return np.sqrt(np.clip(ev, 0.0, None))[::-1].astype(np.float32)

svdvals_topk 

svdvals_topk(A, k: int, oversample: int = 2, power_iters: int = 1, seed: int = 0)

Top-k singular values of a large A, fully on the ANE (sketch + on-ANE qr + projection + on-ANE svd); extra power iters hurt in fp16.

Source code in aneforge/linalg.py 
381
382
383
384
385
386
387
388
389
def svdvals_topk(A, k: int, oversample: int = 2, power_iters: int = 1, seed: int = 0):
  """Top-k singular values of a large A, fully on the ANE (sketch + on-ANE qr + projection + on-ANE svd); extra power iters hurt in fp16."""
  A16 = np.asarray(A, f16); m, n = A16.shape; l = min(k + oversample, n)
  Om = np.random.default_rng(seed).standard_normal((n, l)).astype(f16)
  Y = _ane_gemm(A16, Om); Q = qr(Y)[0].astype(f16)                  # ANE sketch + Gram-Schmidt
  for _ in range(power_iters):
    Y = _ane_gemm(A16, _ane_gemm(A16, Q, transpose_a=True)); Q = qr(Y)[0].astype(f16)
  B = _ane_gemm(A16, Q, transpose_a=True).T                         # ANE projection [l,n]
  return np.sort(svd(B, sweeps=8))[::-1][:k]                         # ANE svd of the small block

qr 

qr(A)

Thin QR A = Q R by modified Gram-Schmidt, fully on the ANE; returns (Q, R).

Source code in aneforge/linalg.py 
409
410
411
412
413
414
def qr(A):
  """Thin QR A = Q R by modified Gram-Schmidt, fully on the ANE; returns (Q, R)."""
  A16 = np.asarray(A, f16); m, n = A16.shape; onem = np.ones((m, 1), f16)
  Qg, Rg = _mgs_qr_graph(af.input((m, n)), m, onem)
  Qv, Rv = _run_multi([Qg, Rg], A16)
  return np.asarray(Qv, np.float32), np.asarray(Rv, np.float32)

cholesky 

cholesky(A)

Lower-triangular Cholesky factor L (A = L L^T) of SPD A, unpivoted, fully on the ANE.

Source code in aneforge/linalg.py 
417
418
419
420
421
422
423
424
425
426
427
428
429
430
def cholesky(A):
  """Lower-triangular Cholesky factor L (A = L L^T) of SPD A, unpivoted, fully on the ANE."""
  A16 = np.asarray(A, f16); n = A16.shape[0]
  At = af.input((n, n)); el = _els_routed(At, n)
  z = el(0, 0) - el(0, 0); L: dict = {}
  for j in range(n):
    d = el(j, j)
    for k in range(j): d = d - L[(j, k)] * L[(j, k)]
    L[(j, j)] = d.relu().adds(1e-12).sqrt()                # relu guards fp16 noise (SPD -> d>0)
    for i in range(j + 1, n):
      s = el(i, j)
      for k in range(j): s = s - L[(i, k)] * L[(j, k)]
      L[(i, j)] = s / L[(j, j)]
  return _solve_once(_grid(L, n, z), A16).astype(np.float32)

lu 

lu(A)

Unpivoted LU (A = L U) by Doolittle, fully on the ANE; returns (L, U).

Source code in aneforge/linalg.py 
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
def lu(A):
  """Unpivoted LU (A = L U) by Doolittle, fully on the ANE; returns (L, U)."""
  A16 = np.asarray(A, f16); n = A16.shape[0]
  At = af.input((n, n)); el = _els_routed(At, n)
  z = el(0, 0) - el(0, 0); one = z.adds(1.0)
  U: dict = {}; L: dict = {(i, i): one for i in range(n)}
  for i in range(n):
    for k in range(i, n):
      s = el(i, k)
      for t in range(i): s = s - L[(i, t)] * U[(t, k)]
      U[(i, k)] = s
    for k in range(i + 1, n):
      s = el(k, i)
      for t in range(i): s = s - L[(k, t)] * U[(t, i)]
      L[(k, i)] = s / U[(i, i)]
  Lv, Uv = _run_multi([_grid(L, n, z), _grid(U, n, z)], A16)
  return np.asarray(Lv, np.float32), np.asarray(Uv, np.float32)

lu_pivoted 

lu_pivoted(A)

Partial-pivoted LU (P A = L U) on the ANE; the per-column argmax pivot is a bridge op, so the program is segmented. Returns (P, L, U).

Source code in aneforge/linalg.py 
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
def lu_pivoted(A):
  """Partial-pivoted LU (P A = L U) on the ANE; the per-column argmax pivot is a bridge op, so the program is segmented. Returns (P, L, U)."""
  A16 = np.asarray(A, f16); n = A16.shape[0]
  At = af.input((n, n)); In = af.input((n, n)); arc = af.input((n, 1))
  arcr = arc.transpose([1, 0]); M = At; P = In
  half = (arc * 0.0).adds(0.5); onerow = (arcr * 0.0).adds(1.0); zerorow = arcr * 0.0
  for k in range(n):
    ek = In.slice_by_size([0, k], [n, 1]); ekr = ek.transpose([1, 0])
    kc = (arc * 0.0).adds(float(k))
    excl = af.select(kc.greater(arc), (arc * 0.0).adds(1e4), arc * 0.0)   # exclude rows < k
    idx = ((M @ ek).abs() - excl).argmax(axis=0)                          # pivot row (bridge op)
    p1h = af.select((arc - idx).abs().greater(half), arc * 0.0, (arc * 0.0).adds(1.0))

    def swap(X):                                                          # swap rows k and pivot
      rk = ekr @ X; rp = p1h.transpose([1, 0]) @ X
      return X + ek @ (rp - rk) + p1h @ (rk - rp)
    M = swap(M); P = swap(P)
    piv = (ekr @ M) @ ek
    belowk = af.select(arc.greater(kc), (arc * 0.0).adds(1.0), arc * 0.0)  # rows > k
    Lk = ((M @ ek) / piv) * belowk
    rowk = ekr @ M
    colabove = af.select(arcr.greater((arcr * 0.0).adds(float(k))), onerow, zerorow)  # cols > k
    M = M - Lk @ (rowk * colabove)                                        # Schur (trailing only)
    M = M - ((M @ ek) * belowk) @ ekr + Lk @ ekr                          # store Lk in column k
  net = af.compile(af.concat([P, M], axis=1))                              # SegmentedModel (argmax cuts)
  o = np.asarray(net(A16, np.eye(n, dtype=f16), np.arange(n, dtype=f16).reshape(n, 1)), np.float32)
  net.release()
  Pv, Mv = o[:, :n], o[:, n:]
  return Pv, np.tril(Mv, -1) + np.eye(n, dtype=np.float32), np.triu(Mv)

generalized_eigh 

generalized_eigh(A, B, sweeps: int = 8)

Eigenvalues of A x = lambda B x (A symmetric, B SPD), ascending, by composing on-engine kernels: B = L L^T, C = L^-1 A L^-T, eigh(C).

Source code in aneforge/linalg.py 
497
498
499
500
501
502
503
504
def generalized_eigh(A, B, sweeps: int = 8):
  """Eigenvalues of A x = lambda B x (A symmetric, B SPD), ascending, by composing on-engine kernels: B = L L^T, C = L^-1 A L^-T, eigh(C)."""
  A16 = np.asarray(A, f16); B16 = np.asarray(B, f16)
  L = cholesky(B16).astype(f16)                          # ANE
  Li = _trinv_lower(L).astype(f16)                       # ANE
  C = _ane_gemm(_ane_gemm(Li, A16), np.ascontiguousarray(Li.T))   # ANE: (L^-1 A) L^-T
  C = ((C + C.T) * 0.5).astype(f16)                      # symmetrize (host, tiny)
  return eigh(C, sweeps=sweeps)

eigvals 

eigvals(A, iters: int = 60)

Eigenvalues of general real A by the unshifted QR algorithm (M <- R Q -> real Schur form), fully on the ANE; returns a complex array.

Source code in aneforge/linalg.py 
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
def eigvals(A, iters: int = 60):
  """Eigenvalues of general real A by the unshifted QR algorithm (M <- R Q -> real Schur form), fully on the ANE; returns a complex array."""
  A16 = np.asarray(A, f16); n = A16.shape[0]; onen = np.ones((n, 1), f16)
  M = af.input((n, n))
  for _ in range(iters):
    Q, R = _mgs_qr_graph(M, n, onen); M = R @ Q        # RQ; unrolls into one program
  Mv = _solve_once(M, A16)
  evs, i = [], 0
  while i < n:
    if i < n - 1 and abs(Mv[i + 1, i]) > 1e-2 * (abs(Mv[i, i]) + abs(Mv[i + 1, i + 1]) + 1e-9):
      b = Mv[i:i + 2, i:i + 2]; tr = b[0, 0] + b[1, 1]
      det = b[0, 0] * b[1, 1] - b[0, 1] * b[1, 0]
      s = np.sqrt(complex(tr * tr - 4 * det)); evs += [(tr + s) / 2, (tr - s) / 2]; i += 2
    else:
      evs.append(complex(Mv[i, i])); i += 1
  return np.array(evs)

solve_triangular 

solve_triangular(A, b, lower: bool = True)

Solve T X = B for triangular T by substitution on the ANE, without forming the inverse (the routed accessor keeps large entries finite, as in cholesky/lu).

b is a vector [n] or a matrix [n, m]; the m columns are solved in one fused graph, since the substitution is a row recurrence and each row is already a full-width slice. A 1-D b returns shape [n], a 2-D one returns [n, m].

Source code in aneforge/linalg.py 
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
def solve_triangular(A, b, lower: bool = True):
  """Solve T X = B for triangular T by substitution on the ANE, without forming the inverse
  (the routed accessor keeps large entries finite, as in cholesky/lu).

  `b` is a vector [n] or a matrix [n, m]; the m columns are solved in one fused graph, since the
  substitution is a row recurrence and each row is already a full-width slice. A 1-D `b` returns
  shape [n], a 2-D one returns [n, m]."""
  A16 = np.asarray(A, f16); n = A16.shape[0]
  b_in = np.asarray(b, f16)
  vec = b_in.ndim == 1
  b16 = b_in.reshape(n, 1) if vec else b_in.reshape(n, -1)
  m = b16.shape[1]
  At = af.input((n, n)); bt = af.input((n, m))
  el = _els_routed(At, n)
  bl = lambda i: bt.slice_by_size([i, 0], [1, m])          # whole row: all m right-hand sides at once
  X: dict = {}
  for i in (range(n) if lower else range(n - 1, -1, -1)):
    s = bl(i)
    for k in (range(i) if lower else range(i + 1, n)): s = s - el(i, k) * X[k]
    X[i] = s / el(i, i)
  out = af.concat([X[i] for i in range(n)], axis=0)
  y = _solve_once(out, A16, b16).reshape(n, m)
  return (y.ravel() if vec else y).astype(np.float32)

solve 

solve(A, b)

Solve A x = b for general square A via the on-ANE pivoted LU (P A = L U).

The direct counterpart to the iterative conjugate_gradient/gmres: factor once, then forward- and back-substitute. b is a vector [n] or a matrix [n, m]. The permutation is applied host-side (P is a permutation matrix, so P b is a row gather, not a matmul worth dispatching).

Source code in aneforge/linalg.py 
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
def solve(A, b):
  """Solve A x = b for general square A via the on-ANE pivoted LU (P A = L U).

  The direct counterpart to the iterative `conjugate_gradient`/`gmres`: factor once, then forward-
  and back-substitute. `b` is a vector [n] or a matrix [n, m]. The permutation is applied host-side
  (P is a permutation matrix, so P b is a row gather, not a matmul worth dispatching)."""
  A16 = np.asarray(A, f16)
  if A16.ndim != 2 or A16.shape[0] != A16.shape[1]:
    raise ValueError(f"solve: A must be square 2-D; got shape {A16.shape}")
  n = A16.shape[0]
  b_in = np.asarray(b, f16)
  if b_in.shape[0] != n:
    raise ValueError(f"solve: b has {b_in.shape[0]} rows but A is {n}x{n}")

  P, L, U = lu_pivoted(A16)
  if np.any(np.diag(U) == 0.0):
    raise np.linalg.LinAlgError("solve: A is singular (zero pivot in U)")
  Pb = np.asarray(P, np.float32) @ np.asarray(b_in, np.float32)          # row permutation of b
  y = solve_triangular(L, Pb.astype(f16), lower=True)
  return solve_triangular(U, np.asarray(y, f16), lower=False)

inv 

inv(A)

A^-1 for general square A, as the matrix solve A X = I on the on-ANE pivoted LU.

One solve against the n identity columns, not n separate solves: solve_triangular already takes a matrix right-hand side and substitutes all columns in one fused graph, so the factorization and both substitutions are dispatched once. Forming the inverse is worse conditioned than solving against the right-hand side you actually have, so prefer solve when you have one.

Source code in aneforge/linalg.py 
628
629
630
631
632
633
634
635
636
637
638
def inv(A):
  """A^-1 for general square A, as the matrix solve A X = I on the on-ANE pivoted LU.

  One `solve` against the n identity columns, not n separate solves: `solve_triangular` already
  takes a matrix right-hand side and substitutes all columns in one fused graph, so the factorization
  and both substitutions are dispatched once. Forming the inverse is worse conditioned than solving
  against the right-hand side you actually have, so prefer `solve` when you have one."""
  A16 = np.asarray(A, f16)
  if A16.ndim != 2 or A16.shape[0] != A16.shape[1]:
    raise ValueError(f"inv: A must be square 2-D; got shape {A16.shape}")
  return solve(A16, np.eye(A16.shape[0], dtype=f16))

lstsq 

lstsq(A, b)

Least-squares solution of the overdetermined A x = b via the on-ANE thin QR.

A = Q R, so the normal equations collapse to R x = Q^T b with R triangular: a direct counterpart to the iterative lsqr, and better conditioned than squaring A into A^T A the way least_squares does. b is a vector [m] or a matrix [m, k]. Requires m >= n and full column rank.

Source code in aneforge/linalg.py 
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
def lstsq(A, b):
  """Least-squares solution of the overdetermined A x = b via the on-ANE thin QR.

  A = Q R, so the normal equations collapse to R x = Q^T b with R triangular: a direct counterpart
  to the iterative `lsqr`, and better conditioned than squaring A into A^T A the way `least_squares`
  does. `b` is a vector [m] or a matrix [m, k]. Requires m >= n and full column rank."""
  A16 = np.asarray(A, f16)
  if A16.ndim != 2:
    raise ValueError(f"lstsq: A must be 2-D; got shape {A16.shape}")
  m, n = A16.shape
  if m < n:
    raise ValueError(f"lstsq: A must be overdetermined (m >= n); got shape {A16.shape}")
  b_in = np.asarray(b, f16)
  if b_in.shape[0] != m:
    raise ValueError(f"lstsq: b has {b_in.shape[0]} rows but A has {m}")

  Q, R = qr(A16)
  if np.any(np.diag(R) == 0.0):
    raise np.linalg.LinAlgError("lstsq: A is rank deficient (zero on the R diagonal)")
  vec = b_in.ndim == 1
  B = np.asarray(b_in, np.float32).reshape(m, 1) if vec else np.asarray(b_in, np.float32)
  Qtb = _ane_gemm(Q.astype(f16), B.astype(f16), transpose_a=True)         # [n, k]
  x = solve_triangular(R, np.asarray(Qtb, f16), lower=False)
  return x.ravel() if vec else x

det 

det(A)

det(A) via the on-ANE pivoted LU (P A = L U): permutation parity times prod(diag U).

Source code in aneforge/linalg.py 
678
679
680
681
682
def det(A):
  """det(A) via the on-ANE pivoted LU (P A = L U): permutation parity times prod(diag U)."""
  P, _, U = lu_pivoted(A)
  sign = _perm_parity(np.argmax(P, axis=1))
  return float(sign * np.prod(np.diag(U).astype(np.float64)))

slogdet 

slogdet(A)

(sign, log|det A|) via the on-ANE pivoted LU; the log-sum form stays finite where the raw product over/underflows. Returns (0.0, -inf) for a singular U diagonal, like numpy.

Source code in aneforge/linalg.py 
685
686
687
688
689
690
691
692
def slogdet(A):
  """(sign, log|det A|) via the on-ANE pivoted LU; the log-sum form stays finite where the raw
  product over/underflows. Returns (0.0, -inf) for a singular U diagonal, like numpy."""
  P, _, U = lu_pivoted(A)
  d = np.diag(U).astype(np.float64)
  if np.any(d == 0.0): return 0.0, float("-inf")
  sign = _perm_parity(np.argmax(P, axis=1)) * float(np.prod(np.sign(d)))
  return float(sign), float(np.sum(np.log(np.abs(d))))

pinv 

pinv(A, k: int, oversample: int = 5, power_iters: int = 2, rcond: float = 0.001)

Rank-k pseudoinverse via randomized_svd: V diag(1/s) U^T with a relative cutoff on small singular values. Rank-k by construction - for a full-rank dense pinv use a dense library.

Source code in aneforge/linalg.py 
695
696
697
698
699
700
701
702
def pinv(A, k: int, oversample: int = 5, power_iters: int = 2, rcond: float = 1e-3):
  """Rank-k pseudoinverse via randomized_svd: V diag(1/s) U^T with a relative cutoff on small
  singular values. Rank-k by construction - for a full-rank dense pinv use a dense library."""
  U, S, Vt = randomized_svd(A, k, oversample=oversample, power_iters=power_iters)
  keep = S > rcond * (S[0] if S.size and S[0] > 0 else 1.0)
  if not np.any(keep): return np.zeros((A.shape[1], A.shape[0]), np.float32)
  W = (Vt[keep].T / S[keep]).astype(np.float32)            # [n, r]
  return _ane_gemm(W.astype(f16), np.ascontiguousarray(U[:, keep].T, np.float32).astype(f16)).astype(np.float32)

norm 

norm(A, order='fro')

Matrix norm of a 2-D A, composed on-engine as one fused graph.

order="fro" is sqrt(sum_square(A)); order=1 is the max absolute column sum and order=inf the max absolute row sum, each an abs, a matmul against a ones vector (the wide accumulator does the summing), and an amax. Oracle is np.linalg.norm (whose argument is spelled ord; renamed here only to avoid shadowing the builtin).

Source code in aneforge/linalg.py 
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
def norm(A, order="fro"):
  """Matrix norm of a 2-D A, composed on-engine as one fused graph.

  `order="fro"` is `sqrt(sum_square(A))`; `order=1` is the max absolute column sum and
  `order=inf` the max absolute row sum, each an abs, a matmul against a ones vector (the
  wide accumulator does the summing), and an `amax`. Oracle is `np.linalg.norm` (whose
  argument is spelled `ord`; renamed here only to avoid shadowing the builtin).
  """
  A16 = np.asarray(A, f16)
  if A16.ndim != 2: raise ValueError(f"linalg.norm: expected a 2-D matrix; got shape {A16.shape}")
  m, n = A16.shape
  X = af.input((m, n))
  if order == "fro":
    out = X.sum_square((0, 1)).sqrt()
  elif order == 1:
    # column sums: |A|^T @ ones[m,1] -> [n,1]
    out = (X.abs().transpose([1, 0]) @ np.ones((m, 1), f16)).amax((0, 1))
  elif order in (np.inf, float("inf")):
    out = (X.abs() @ np.ones((n, 1), f16)).amax((0, 1))     # row sums -> [m,1]
  else:
    raise ValueError(f"linalg.norm: order must be 'fro', 1, or inf; got {order!r}")
  net = af.compile(out, _check_precision=False)
  v = float(np.ravel(np.asarray(net(A16), np.float64))[0])
  net.release()
  return v

kron 

kron(A, B)

Kronecker product A (x) B by broadcast-multiply of expanded views; result [mp, nq].

Source code in aneforge/linalg.py 
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
def kron(A, B):
  """Kronecker product A (x) B by broadcast-multiply of expanded views; result [m*p, n*q]."""
  A16 = np.asarray(A, f16); B16 = np.asarray(B, f16)
  if A16.ndim != 2 or B16.ndim != 2:
    raise ValueError(f"linalg.kron: expected 2-D matrices; got {A16.shape} and {B16.shape}")
  m, n = A16.shape; p, q = B16.shape
  At = af.input((m, n)); Bt = af.input((p, q))
  A4 = At.expand_dims((1, 3))
  B4 = Bt.expand_dims((0, 2))
  C = A4 * B4
  out = C.reshape(m * p, n * q)
  net = af.compile(out, _check_precision=False)
  Y = net(A16, B16)
  net.release()
  return np.asarray(Y, np.float32)

expm 

expm(A, order: int = 8)

exp(A) by scaling and squaring: exp(A) = exp(A / 2^s)^(2^s), with the inner exponential a Taylor sum of order terms.

The host picks s from norm(A, 1) so the scaled matrix has norm <= 1/2, which is where a low-order Taylor sum is accurate; everything after that is on-engine gemms, O(order + s) of them.

fp16 squaring compounds error the same way matrix_power does, so this is for modest norms on well-conditioned matrices. Oracle is scipy.linalg.expm.

Source code in aneforge/linalg.py 
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
def expm(A, order: int = 8):
  """exp(A) by scaling and squaring: exp(A) = exp(A / 2^s)^(2^s), with the inner exponential a
  Taylor sum of `order` terms.

  The host picks `s` from `norm(A, 1)` so the scaled matrix has norm <= 1/2, which is where a
  low-order Taylor sum is accurate; everything after that is on-engine gemms, O(order + s) of them.

  fp16 squaring compounds error the same way `matrix_power` does, so this is for modest norms on
  well-conditioned matrices. Oracle is `scipy.linalg.expm`.
  """
  A16 = np.asarray(A, f16)
  if A16.ndim != 2 or A16.shape[0] != A16.shape[1]:
    raise ValueError(f"linalg.expm: expected a square matrix; got shape {A16.shape}")
  if order < 1:
    raise ValueError(f"linalg.expm: order must be >= 1; got {order}")
  n = A16.shape[0]
  nrm = float(np.abs(np.asarray(A16, np.float64)).sum(axis=0).max())      # 1-norm, max abs col sum
  s = max(0, int(np.ceil(np.log2(nrm / 0.5)))) if nrm > 0.5 else 0
  As = (np.asarray(A16, np.float64) / (2.0 ** s)).astype(f16)

  # Taylor: I + As + As^2/2! + ... The term recurrence keeps one gemm per order rather than
  # recomputing powers, and 1/k! folds into the host-side scalar so no extra dispatch.
  acc = np.eye(n, dtype=np.float64)
  term = np.eye(n, dtype=np.float64)
  for k in range(1, order + 1):
    term = np.asarray(_ane_gemm(term.astype(f16), As), np.float64) / k
    acc = acc + term
  for _ in range(s):                                                     # undo the scaling
    acc = np.asarray(_ane_gemm(acc.astype(f16), acc.astype(f16)), np.float64)
  return acc.astype(np.float32)

matrix_power 

matrix_power(A, n: int)

A**n for integer n >= 0 by binary exponentiation over _ane_gemm, so it costs O(log n) on-engine gemms rather than n-1. n=0 is the identity.

fp16 error compounds with each squaring, so this is for modest powers on well-conditioned matrices; negative n needs an inverse and is rejected. Oracle is np.linalg.matrix_power.

Source code in aneforge/linalg.py 
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
def matrix_power(A, n: int):
  """A**n for integer n >= 0 by binary exponentiation over `_ane_gemm`, so it costs
  O(log n) on-engine gemms rather than n-1. n=0 is the identity.

  fp16 error compounds with each squaring, so this is for modest powers on
  well-conditioned matrices; negative n needs an inverse and is rejected.
  Oracle is `np.linalg.matrix_power`.
  """
  A16 = np.asarray(A, f16)
  if A16.ndim != 2 or A16.shape[0] != A16.shape[1]:
    raise ValueError(f"linalg.matrix_power: expected a square matrix; got shape {A16.shape}")
  if n < 0:
    raise ValueError("linalg.matrix_power: negative powers need a matrix inverse; not supported")
  if n == 0: return np.eye(A16.shape[0], dtype=np.float32)
  acc, base = None, A16                                     # square-and-multiply
  while n:
    if n & 1: acc = base.copy() if acc is None else _ane_gemm(acc, base)
    n >>= 1
    if n: base = _ane_gemm(base, base)
  return np.asarray(acc, np.float32)

polar 

polar(A)

Polar decomposition A = U @ P, on the ANE.

U is orthogonal for square A and semi-orthogonal for either rectangular shape (U^T U = I for tall A, U U^T = I for wide A); P is symmetric positive-semidefinite [n,n]. Composed from the on-ANE SVD: U_ S V^T -> U = U_ V^T, P = V diag(S) V^T. Oracle: scipy.linalg.polar(A, side='right').

Source code in aneforge/linalg.py 
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
def polar(A):
  """Polar decomposition A = U @ P, on the ANE.

  U is orthogonal for square A and semi-orthogonal for either rectangular shape (U^T U = I for
  tall A, U U^T = I for wide A); P is symmetric positive-semidefinite [n,n].
  Composed from the on-ANE SVD: U_ S V^T -> U = U_ V^T, P = V diag(S) V^T.
  Oracle: scipy.linalg.polar(A, side='right')."""
  A16 = np.asarray(A, f16)
  if A16.ndim != 2: raise ValueError(f"polar: expected 2-D; got shape {A16.shape}")
  m, n = A16.shape
  U_svd, S, Vt = randomized_svd(A16, k=min(m, n), oversample=5, power_iters=2)
  V = Vt.T.astype(f16)
  Umat = _ane_gemm(U_svd.astype(f16), Vt)                        # U_ @ V^T -> [m,n]
  Sdiag = np.diag(S).astype(f16)
  Pmat = _ane_gemm(_ane_gemm(V, Sdiag), Vt)                      # V diag(S) V^T -> [n,n]
  P = np.asarray(Pmat, np.float32)
  P = 0.5 * (P + P.T)                                            # fp16 GEMMs leave ~1e-4 asymmetry; sym kills it
  return np.asarray(Umat, np.float32), P

matrix_rank 

matrix_rank(A, tol=None)

Numerical rank of A by counting singular values above tol, on the ANE.

Default tol is max(m, n) * 1e-5 * sigma_max, mirroring numpy's max(m, n) * eps * sigma_max but with the measured floor of this backend in place of a machine epsilon -- the input is fp16, so nothing here resolves a singular value below ~1e-5 of sigma_max. Pass tol explicitly if you know your spectrum. Returns 0 for a zero matrix. Oracle: np.linalg.matrix_rank.

Source code in aneforge/linalg.py 
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
def matrix_rank(A, tol=None):
  """Numerical rank of A by counting singular values above `tol`, on the ANE.

  Default `tol` is max(m, n) * 1e-5 * sigma_max, mirroring numpy's max(m, n) * eps * sigma_max but
  with the measured floor of this backend in place of a machine epsilon -- the input is fp16, so
  nothing here resolves a singular value below ~1e-5 of sigma_max.  Pass `tol` explicitly if you
  know your spectrum.  Returns 0 for a zero matrix.  Oracle: np.linalg.matrix_rank."""
  A16 = np.asarray(A, f16)
  if A16.ndim != 2: raise ValueError(f"matrix_rank: expected 2-D; got shape {A16.shape}")
  m, n = A16.shape
  S = _svals(A16)
  if S.size == 0: return 0
  if tol is None:
    tol = max(m, n) * _RANK_FLOOR * float(S[0])
  return int(np.sum(S > tol))

cond 

cond(A)

2-norm condition number sigma_max / sigma_min via SVD, on the ANE.

Works for square and rectangular A. Accurate to cond ~1e3, past which the fp16 input itself stops resolving sigma_min. Returns inf for a zero or exactly singular matrix. Oracle: np.linalg.cond(A).

Source code in aneforge/linalg.py 
854
855
856
857
858
859
860
861
862
863
864
def cond(A):
  """2-norm condition number sigma_max / sigma_min via SVD, on the ANE.

  Works for square and rectangular A.  Accurate to cond ~1e3, past which the fp16 input itself
  stops resolving sigma_min.  Returns inf for a zero or exactly singular matrix.
  Oracle: np.linalg.cond(A)."""
  A16 = np.asarray(A, f16)
  if A16.ndim != 2: raise ValueError(f"cond: expected 2-D; got shape {A16.shape}")
  S = _svals(A16)
  if S.size == 0 or float(S[-1]) <= 0.0: return float("inf")
  return float(S[0]) / float(S[-1])

FFT

fft

aneforge.fft - staged Cooley-Tukey FFT on the ANE (complex carried as (re, im) pairs, each stage a matmul).

Plan

A compiled staged-FFT program plus the cross-twiddle constants it threads in each call.

Source code in aneforge/fft.py 
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
class Plan:
  """A compiled staged-FFT program plus the cross-twiddle constants it threads in each call."""

  def __init__(self, N: int, inverse: bool, real_input: bool):
    self.N = N
    self.inverse = inverse
    self.real_input = real_input
    b = _Builder(inverse)
    # real/imag inputs; for rfft the imag input is fed as zeros.
    xr = af.input((1, N)); xi = af.input((1, N))
    Xr, Xi = b.transform(xr, xi, N)
    if inverse:
      Xr = Xr * (1.0 / N)
      Xi = Xi * (1.0 / N)
    out = af.concat([Xr, Xi], axis=1)        # [1, 2N], split on host
    self._aux_values = b.aux_values
    self.n_stages = _stage_count(N)          # number of dense-DFT matmul stages (<=3)
    self.model = af.compile(out, _check_precision=False)
    self.n_ops = self.model.n_ops

  def __call__(self, x_re: np.ndarray, x_im: np.ndarray | None = None):
    x_re = np.asarray(x_re, np.float16).reshape(1, self.N)
    if x_im is None:
      x_im = np.zeros((1, self.N), np.float16)
    else:
      x_im = np.asarray(x_im, np.float16).reshape(1, self.N)
    out = self.model(x_re, x_im, *self._aux_values)
    out = out.reshape(2, self.N)
    return out[0].copy(), out[1].copy()

Plan2

A compiled 2-D FFT for [M,N] complex fields, one fused program (separable F_M @ X @ F_N^T, eight real GEMMs).

Source code in aneforge/fft.py 
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
class Plan2:
  """A compiled 2-D FFT for [M,N] complex fields, one fused program (separable F_M @ X @ F_N^T, eight real GEMMs)."""

  def __init__(self, M: int, N: int, inverse: bool):
    self.M, self.N, self.inverse = M, N, inverse
    mk = _idft_matrix if inverse else _dft_matrix
    WrN, WiN = mk(N)                                          # row twiddle (x @ Wt)
    WrM, WiM = mk(M)                                          # column twiddle
    if inverse:
        # fold 1/(M*N) into the twiddles per-pass (an unscaled first-axis transform overflows fp16).
      WrN, WiN = WrN * (1.0 / N), WiN * (1.0 / N)
      WrM, WiM = WrM * (1.0 / M), WiM * (1.0 / M)
    xr = af.input((M, N)); xi = af.input((M, N))
    re, im = _cmatmul_const(xr, xi, WrN, WiN)                 # all M rows, one matmul
    re = re.transpose([1, 0]); im = im.transpose([1, 0])      # columns -> rows
    re, im = _cmatmul_const(re, im, WrM, WiM)                 # all N columns, one matmul
    re = re.transpose([1, 0]); im = im.transpose([1, 0])
    out = af.concat([re, im], axis=0)                          # [2M, N], split on host
    self.model = af.compile(out, _check_precision=False)
    self.n_ops = self.model.n_ops

  def __call__(self, x_re: np.ndarray, x_im: np.ndarray | None = None):
    x_re = np.asarray(x_re, np.float16).reshape(self.M, self.N)
    if x_im is None:
      x_im = np.zeros((self.M, self.N), np.float16)
    else:
      x_im = np.asarray(x_im, np.float16).reshape(self.M, self.N)
    out = self.model(x_re, x_im).reshape(2, self.M, self.N)
    return out[0].copy(), out[1].copy()

fft 

fft(x_re, x_im, N: int)

Forward FFT of a complex signal (real/imag arrays), length N, on the ANE; returns (X_re, X_im).

Source code in aneforge/fft.py 
277
278
279
def fft(x_re, x_im, N: int):
  """Forward FFT of a complex signal (real/imag arrays), length N, on the ANE; returns (X_re, X_im)."""
  return fft_plan(N)(x_re, x_im)

ifft 

ifft(X_re, X_im, N: int)

Inverse FFT (1/N normalized) of a complex spectrum on the ANE; returns (x_re, x_im).

Source code in aneforge/fft.py 
282
283
284
def ifft(X_re, X_im, N: int):
  """Inverse FFT (1/N normalized) of a complex spectrum on the ANE; returns (x_re, x_im)."""
  return ifft_plan(N)(X_re, X_im)

rfft 

rfft(x_real, N: int)

Forward FFT of a real signal (imag = 0) on the ANE; returns the full-length spectrum (X_re, X_im).

Source code in aneforge/fft.py 
287
288
289
def rfft(x_real, N: int):
  """Forward FFT of a real signal (imag = 0) on the ANE; returns the full-length spectrum (X_re, X_im)."""
  return rfft_plan(N)(x_real, None)

irfft 

irfft(X_re, X_im, N: int)

Inverse real FFT of a Hermitian-symmetric spectrum on the ANE; returns the real time-domain signal of length N (the imag part is ~0 by Hermitian symmetry, and numpy.fft.irfft also returns only the real part).

Takes the FULL length-N spectrum, as rfft returns it -- not the N//2+1 half spectrum that numpy.fft.irfft expects.

Source code in aneforge/fft.py 
292
293
294
295
296
297
298
299
300
def irfft(X_re, X_im, N: int):
  """Inverse real FFT of a Hermitian-symmetric spectrum on the ANE; returns the real time-domain
  signal of length N (the imag part is ~0 by Hermitian symmetry, and numpy.fft.irfft also
  returns only the real part).

  Takes the FULL length-N spectrum, as `rfft` returns it -- not the N//2+1 half spectrum that
  numpy.fft.irfft expects."""
  x_re, _ = ifft_plan(N)(X_re, X_im)
  return x_re

fft2 

fft2(x_re, x_im=None)

2-D FFT of an [M,N] complex field on the ANE (x_im=None means a real field); returns (X_re, X_im).

Source code in aneforge/fft.py 
303
304
305
306
307
def fft2(x_re, x_im=None):
  """2-D FFT of an [M,N] complex field on the ANE (x_im=None means a real field); returns (X_re, X_im)."""
  x_re = np.asarray(x_re)
  M, N = x_re.shape
  return fft2_plan(M, N)(x_re, x_im)

ifft2 

ifft2(X_re, X_im)

Inverse 2-D FFT (1/(M*N) normalized) on the ANE; returns (x_re, x_im).

Source code in aneforge/fft.py 
310
311
312
313
314
def ifft2(X_re, X_im):
  """Inverse 2-D FFT (1/(M*N) normalized) on the ANE; returns (x_re, x_im)."""
  X_re = np.asarray(X_re)
  M, N = X_re.shape
  return ifft2_plan(M, N)(X_re, X_im)

magnitude 

magnitude(X_re, X_im)

|X| = sqrt(re^2 + im^2) (numpy on host outputs).

Source code in aneforge/fft.py 
317
318
319
def magnitude(X_re, X_im):
  """|X| = sqrt(re^2 + im^2) (numpy on host outputs)."""
  return np.sqrt(np.asarray(X_re, np.float32) ** 2 + np.asarray(X_im, np.float32) ** 2)

power 

power(X_re, X_im)

|X|^2 power spectrum (numpy on host outputs).

Source code in aneforge/fft.py 
322
323
324
def power(X_re, X_im):
  """|X|^2 power spectrum (numpy on host outputs)."""
  return np.asarray(X_re, np.float32) ** 2 + np.asarray(X_im, np.float32) ** 2

Special functions

special

aneforge.special - special functions as fused fp16 polynomial chains on the ANE; each takes and returns an aneforge.Tensor.

sin 

sin(x: Tensor) -> Tensor

sin(x) for x in [-pi/2, pi/2], portable fp16 polynomial; reduce wider args on the host.

Source code in aneforge/special.py 
40
41
42
def sin(x: Tensor) -> Tensor:
  """sin(x) for x in [-pi/2, pi/2], portable fp16 polynomial; reduce wider args on the host."""
  return x * _poly_in(x * x, _SIN_P)

cos 

cos(x: Tensor) -> Tensor

cos(x) for x in [-pi/2, pi/2], portable fp16 polynomial; reduce wider args on the host.

Source code in aneforge/special.py 
45
46
47
def cos(x: Tensor) -> Tensor:
  """cos(x) for x in [-pi/2, pi/2], portable fp16 polynomial; reduce wider args on the host."""
  return _poly_in(x * x, _COS_Q)

erfc 

erfc(x: Tensor) -> Tensor

Complementary error function erfc(x) for x in [0, ~6] (direct A&S form, no cancellation for large x).

Source code in aneforge/special.py 
57
58
59
60
61
def erfc(x: Tensor) -> Tensor:
  """Complementary error function erfc(x) for x in [0, ~6] (direct A&S form, no cancellation for large x)."""
  t = _const(x, 1.0) / (_const(x, 1.0) + x * _ERFC_P)
  poly = _poly_in(t, _ERFC_A) * t                         # a0..a4 low-first, then *t
  return poly * (x * x * -1.0).exp()

erf 

erf(x: Tensor) -> Tensor

Error function erf(x) for |x| <= 2, as x*poly(x^2) (deg-5 minimax of erf(x)/x); odd, so it holds for negative x. Past |x| ~ 2 use 1 - erfc(|x|), where erfc is small and the subtraction no longer cancels.

Source code in aneforge/special.py 
73
74
75
def erf(x: Tensor) -> Tensor:
  """Error function erf(x) for |x| <= 2, as x*poly(x^2) (deg-5 minimax of erf(x)/x); odd, so it holds for negative x. Past |x| ~ 2 use 1 - erfc(|x|), where erfc is small and the subtraction no longer cancels."""
  return x * _poly_in(x * x, _ERF_P)

expm1 

expm1(x: Tensor) -> Tensor

exp(x) - 1 accurate near 0, Taylor for |x| <= ~0.7.

Source code in aneforge/special.py 
80
81
82
83
def expm1(x: Tensor) -> Tensor:
  """exp(x) - 1 accurate near 0, Taylor for |x| <= ~0.7."""
  inner = _horner(x, [1.0 / 720, 1.0 / 120, 1.0 / 24, 1.0 / 6, 0.5, 1.0])
  return x * inner

log1p 

log1p(x: Tensor) -> Tensor

log(1 + x) accurate near 0, as x * poly(x) (deg-7 minimax of log1p(x)/x) for x in [-0.5, 1.0].

Source code in aneforge/special.py 
91
92
93
def log1p(x: Tensor) -> Tensor:
  """log(1 + x) accurate near 0, as x * poly(x) (deg-7 minimax of log1p(x)/x) for x in [-0.5, 1.0]."""
  return x * _horner(x, _LOG1P_RATIO)

lgamma 

lgamma(x: Tensor) -> Tensor

Log-gamma log|Gamma(x)| for x in [1, 8], deg-8 minimax centered at x-4.5 (accurate in absolute terms; rel error ill-defined at the zeros x=1,2).

Source code in aneforge/special.py 
111
112
113
def lgamma(x: Tensor) -> Tensor:
  """Log-gamma log|Gamma(x)| for x in [1, 8], deg-8 minimax centered at x-4.5 (accurate in absolute terms; rel error ill-defined at the zeros x=1,2)."""
  return _horner(x + _const(x, -_LGAMMA_C), _LGAMMA)

gamma 

gamma(x: Tensor) -> Tensor

Gamma function on x in [1, 2], deg-6 minimax centered at x-1.5; fp16-narrow (use gamma_via_lgamma for a wider range).

Source code in aneforge/special.py 
116
117
118
def gamma(x: Tensor) -> Tensor:
  """Gamma function on x in [1, 2], deg-6 minimax centered at x-1.5; fp16-narrow (use gamma_via_lgamma for a wider range)."""
  return _horner(x + _const(x, -_GAMMA_C), _GAMMA_12)

gamma_via_lgamma 

gamma_via_lgamma(x: Tensor) -> Tensor

Gamma(x) = exp(lgamma(x)) for x in [1, ~7.5]; wider range than gamma at a small accuracy cost.

Source code in aneforge/special.py 
121
122
123
def gamma_via_lgamma(x: Tensor) -> Tensor:
  """Gamma(x) = exp(lgamma(x)) for x in [1, ~7.5]; wider range than `gamma` at a small accuracy cost."""
  return lgamma(x).exp()

bessel_j0 

bessel_j0(x: Tensor) -> Tensor

Bessel J0(x) for |x| <= 3 (A&S 9.4.1 in (x/3)^2).

Source code in aneforge/special.py 
136
137
138
139
def bessel_j0(x: Tensor) -> Tensor:
  """Bessel J0(x) for |x| <= 3 (A&S 9.4.1 in (x/3)^2)."""
  t = (x * x) * (1.0 / 9.0)
  return _poly_in(t, _J0)

bessel_i0 

bessel_i0(x: Tensor) -> Tensor

Modified Bessel I0(x) for |x| <= 3.75 (A&S 9.8.1); overflows fp16 past x~12.

Source code in aneforge/special.py 
142
143
144
145
def bessel_i0(x: Tensor) -> Tensor:
  """Modified Bessel I0(x) for |x| <= 3.75 (A&S 9.8.1); overflows fp16 past x~12."""
  t = (x * x) * (1.0 / (3.75 * 3.75))
  return _poly_in(t, _I0)

bessel_k0 

bessel_k0(x: Tensor) -> Tensor

Modified Bessel K0(x) for 0 < x <= 2 (A&S 9.8.5): -ln(x/2) I0(x) + series((x/2)^2).

Source code in aneforge/special.py 
148
149
150
151
152
153
154
def bessel_k0(x: Tensor) -> Tensor:
  """Modified Bessel K0(x) for 0 < x <= 2 (A&S 9.8.5): -ln(x/2) I0(x) + series((x/2)^2)."""
  half = x * 0.5
  t_k0 = half * half                       # (x/2)^2  for the K0 series
  t_i0 = (x * x) * (1.0 / (3.75 * 3.75))    # (x/3.75)^2 for the I0 factor
  i0 = _poly_in(t_i0, _I0)
  return half.log() * (i0 * -1.0) + _poly_in(t_k0, _K0)

bessel_j1 

bessel_j1(x: Tensor) -> Tensor

Bessel J1(x) for |x| <= 3 (A&S 9.4.3): x/2 * poly((x/3)^2).

Source code in aneforge/special.py 
163
164
165
166
def bessel_j1(x: Tensor) -> Tensor:
  """Bessel J1(x) for |x| <= 3 (A&S 9.4.3): x/2 * poly((x/3)^2)."""
  t = (x * x) * (1.0 / 9.0)
  return x * _const(x, 0.5) * _poly_in(t, _J1)

bessel_i1 

bessel_i1(x: Tensor) -> Tensor

Modified Bessel I1(x) for |x| <= 3.75 (A&S 9.8.3): x/2 * poly((x/3.75)^2). Overflows fp16 past x~12.

Source code in aneforge/special.py 
169
170
171
172
def bessel_i1(x: Tensor) -> Tensor:
  """Modified Bessel I1(x) for |x| <= 3.75 (A&S 9.8.3): x/2 * poly((x/3.75)^2). Overflows fp16 past x~12."""
  t = (x * x) * (1.0 / (3.75 * 3.75))
  return x * _const(x, 0.5) * _poly_in(t, _I1)

digamma 

digamma(x: Tensor) -> Tensor

Digamma (psi) function for x > 0. Shifts by one via the exact recurrence psi(x) = psi(x+1) - 1/x, then applies the asymptotic expansion log(x+1) - 1/(2(x+1)) - B2/(2(x+1)^2) - B4/(4(x+1)^4) - B6/(6*(x+1)^6). For 0 < x < 1 the -1/x term dominates and the result is still accurate.

Source code in aneforge/special.py 
187
188
189
190
191
192
193
194
195
def digamma(x: Tensor) -> Tensor:
  """Digamma (psi) function for x > 0. Shifts by one via the exact recurrence psi(x) = psi(x+1) - 1/x, then applies the asymptotic expansion log(x+1) - 1/(2(x+1)) - B2/(2*(x+1)^2) - B4/(4*(x+1)^4) - B6/(6*(x+1)^6). For 0 < x < 1 the -1/x term dominates and the result is still accurate."""
  # shift x by one via psi(x) = psi(x+1) - 1/x
  x2 = x + _const(x, 1.0)       # x+1; for x >= 1 this puts x2 in the asymptotic range
  shift = _const(x, 1.0) / x    # 1/x term from the recurrence
  # asymptotic expansion at x2 = x+1
  inv = _const(x, 1.0) / x2
  inv2 = inv * inv
  return x2.log() - inv * _const(x, 0.5) - inv2 * (_const(x, _DIGAMMA_B2_12) + inv2 * (_const(x, -_DIGAMMA_B4_120) + inv2 * _const(x, _DIGAMMA_B6_252))) - shift

beta 

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

Beta(a, b) = exp(lgamma(a) + lgamma(b) - lgamma(a + b)), for a, b > 0. Most accurate when a, b >= 2.5 (away from lgamma's zeros at 1, 2); near those zeros the absolute error in lgamma (~0.2) propagates through the exp. Avoids the overflow of gamma(a)*gamma(b)/gamma(a+b).

Source code in aneforge/special.py 
200
201
202
def beta(a: Tensor, b: Tensor) -> Tensor:
  """Beta(a, b) = exp(lgamma(a) + lgamma(b) - lgamma(a + b)), for a, b > 0. Most accurate when a, b >= 2.5 (away from lgamma's zeros at 1, 2); near those zeros the absolute error in lgamma (~0.2) propagates through the exp. Avoids the overflow of gamma(a)*gamma(b)/gamma(a+b)."""
  return (lgamma(a) + lgamma(b) - lgamma(a + b)).exp()

exp_wide 

exp_wide(x: Tensor, splits: int = 1) -> Tensor

exp for wide x via repeated squaring; does NOT reliably beat the native x.exp() (documented, not a win).

Source code in aneforge/special.py 
207
208
209
210
211
def exp_wide(x: Tensor, splits: int = 1) -> Tensor:
  """exp for wide `x` via repeated squaring; does NOT reliably beat the native x.exp() (documented, not a win)."""
  e = (x * (1.0 / (1 << splits))).exp()
  for _ in range(splits): e = e * e
  return e

log_wide 

log_wide(x: Tensor, sqrts: int = 3) -> Tensor

log for wide positive x via repeated sqrt + scale-back; matches rather than beats the native x.log().

Source code in aneforge/special.py 
214
215
216
217
218
def log_wide(x: Tensor, sqrts: int = 3) -> Tensor:
  """log for wide positive `x` via repeated sqrt + scale-back; matches rather than beats the native x.log()."""
  r = x
  for _ in range(sqrts): r = r.sqrt()
  return r.log() * float(1 << sqrts)

sinh 

sinh(x: Tensor) -> Tensor

sinh(x) for |x| <= ~10; overflows fp16 near ln(65504) ~ 11 (same wall as softplus).

Accurate in absolute terms, not relative: e^x - e^-x cancels near the origin, so the relative error reaches ~2.4% for |x| < 0.01 (absolute error stays ~2e-5). Use expm1 if you need the small-argument regime; its Taylor form is only valid for |x| <= ~0.7.

Source code in aneforge/special.py 
223
224
225
226
227
228
229
def sinh(x: Tensor) -> Tensor:
  """sinh(x) for |x| <= ~10; overflows fp16 near ln(65504) ~ 11 (same wall as softplus).

  Accurate in absolute terms, not relative: e^x - e^-x cancels near the origin, so the relative
  error reaches ~2.4% for |x| < 0.01 (absolute error stays ~2e-5). Use expm1 if you need the
  small-argument regime; its Taylor form is only valid for |x| <= ~0.7."""
  return (x.exp() - (x * -1.0).exp()) * 0.5

cosh 

cosh(x: Tensor) -> Tensor

cosh(x) for |x| <= ~10; overflows fp16 near ln(65504) ~ 11 (same wall as softplus).

Source code in aneforge/special.py 
232
233
234
def cosh(x: Tensor) -> Tensor:
  """cosh(x) for |x| <= ~10; overflows fp16 near ln(65504) ~ 11 (same wall as softplus)."""
  return (x.exp() + (x * -1.0).exp()) * 0.5

asinh 

asinh(x: Tensor) -> Tensor

asinh(x) for |x| <= ~10; domain is all real x, but fp16 overflows for large |x|.

Evaluated on |x| and signed back: asinh is odd, and log(x + sqrt(x^2+1)) cancels for x << 0 (2.1% relative error at x=-10 in fp16, vs 0.0% for this form).

Source code in aneforge/special.py 
237
238
239
240
241
242
def asinh(x: Tensor) -> Tensor:
  """asinh(x) for |x| <= ~10; domain is all real x, but fp16 overflows for large |x|.

  Evaluated on |x| and signed back: asinh is odd, and log(x + sqrt(x^2+1)) cancels for x << 0
  (2.1% relative error at x=-10 in fp16, vs 0.0% for this form)."""
  return (x.abs() + (x * x).adds(1.0).sqrt()).log() * x.sign()

acosh 

acosh(x: Tensor) -> Tensor

acosh(x) for x >= 1; fp16 overflows for large x.

Source code in aneforge/special.py 
245
246
247
def acosh(x: Tensor) -> Tensor:
  """acosh(x) for x >= 1; fp16 overflows for large x."""
  return (x + (x * x).adds(-1.0).sqrt()).log()

atanh 

atanh(x: Tensor) -> Tensor

atanh(x) for |x| < 1; singular at |x| == 1.

Accurate in absolute terms, not relative: (1+x)/(1-x) rounds to ~1 near the origin, so the relative error reaches ~4% for |x| < 0.01 (absolute error stays ~1e-4). log1p fixes the small-argument regime but is only valid for its argument in [-0.5, 1], i.e. |x| <= 1/3.

Source code in aneforge/special.py 
250
251
252
253
254
255
256
def atanh(x: Tensor) -> Tensor:
  """atanh(x) for |x| < 1; singular at |x| == 1.

  Accurate in absolute terms, not relative: (1+x)/(1-x) rounds to ~1 near the origin, so the
  relative error reaches ~4% for |x| < 0.01 (absolute error stays ~1e-4). log1p fixes the
  small-argument regime but is only valid for its argument in [-0.5, 1], i.e. |x| <= 1/3."""
  return ((x.adds(1.0)) / (x * -1.0).adds(1.0)).log() * 0.5

einsum

einsum

aneforge.einsum - numpy-style einsum lowered to aneforge graph ops as one fused ANE program.

EinsumUnsupported

Bases: NotImplementedError

Raised for equations not lowerable to ANE matmul/reduce ops (diagonal-write outputs, triple-repeat indices, and ellipsis broadcast between mismatched size-1 dims). Plain ellipsis batch dims are supported.

Source code in aneforge/einsum.py 
14
15
16
class EinsumUnsupported(NotImplementedError):
  """Raised for equations not lowerable to ANE matmul/reduce ops (diagonal-write outputs, triple-repeat
  indices, and ellipsis broadcast between mismatched size-1 dims). Plain ellipsis batch dims are supported."""

einsum 

einsum(equation: str, *operands: Tensor) -> Tensor

numpy-style einsum lowered to aneforge ops (matmul-reducible patterns, incl. '...' batch dims and diagonal/trace extraction); rejects diagonal-WRITE and repeated output indices.

Source code in aneforge/einsum.py 
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
def einsum(equation: str, *operands: Tensor) -> Tensor:
  """numpy-style einsum lowered to aneforge ops (matmul-reducible patterns, incl. '...' batch
  dims and diagonal/trace extraction); rejects diagonal-WRITE and repeated output indices."""
  if not operands:
    raise ValueError("einsum: at least one operand is required")
  if not all(isinstance(o, Tensor) for o in operands):
    raise TypeError("einsum: all operands must be aneforge Tensors (build them with af.input)")

  equation = _expand_ellipsis(equation, operands)
  ins, out = _parse(equation, len(operands))
  for s, t in zip(ins, operands):
    if len(s) != len(t.shape):
      raise ValueError(f"einsum: subscript '{s}' has {len(s)} indices but operand "
               f"shape {t.shape} has {len(t.shape)} dims")

  # a doubly-repeated index within an operand is a diagonal: take it before contracting, so
  # everything downstream sees operands whose indices are distinct.
  operands, ins = zip(*(_extract_diagonals(t, s) for t, s in zip(operands, ins)))
  ins = list(ins)

  # single operand: pure transpose / reduce
  if len(operands) == 1:
    t, sub = operands[0], ins[0]
    t, sub = _sum_to(t, sub, out)            # drop summed indices
    if sub == out: return t
    return _align(t, sub, out)               # transpose to requested order

  # multi-operand: left-fold pairwise
  cur, csub = operands[0], ins[0]
  for i in range(1, len(operands)):
    nxt, nsub = operands[i], ins[i]
    remaining = "".join(ins[i + 1:])
    # keep anything in the final output or still needed by a later operand.
    survive = set(out) | set(remaining)
    step_out = "".join(dict.fromkeys(
      [ch for ch in csub + nsub if ch in survive]))
    cur, csub = _contract_pair(cur, csub, nxt, nsub, step_out)

  # final: sum away any leftover non-output indices, then order to `out`
  cur, csub = _sum_to(cur, csub, out)
  if csub != out: cur = _align(cur, csub, out)
  return cur

Signal processing

dsp

aneforge.dsp - DSP toolkit on the ANE, composed from conv and the matmul-FFT (aneforge.fft).

hann 

hann(M: int, sym: bool = False) -> np.ndarray

Hann window (raised cosine); sym=False gives the periodic/DFT (STFT default) form.

Source code in aneforge/dsp.py 
21
22
23
24
25
26
def hann(M: int, sym: bool = False) -> np.ndarray:
  """Hann window (raised cosine); `sym=False` gives the periodic/DFT (STFT default) form."""
  if M == 1: return np.ones(1, np.float32)
  n = np.arange(M)
  denom = (M - 1) if sym else M
  return (0.5 - 0.5 * np.cos(2.0 * np.pi * n / denom)).astype(np.float32)

hamming 

hamming(M: int, sym: bool = False) -> np.ndarray

Hamming window. sym=False = periodic (STFT) form.

Source code in aneforge/dsp.py 
29
30
31
32
33
34
def hamming(M: int, sym: bool = False) -> np.ndarray:
  """Hamming window. `sym=False` = periodic (STFT) form."""
  if M == 1: return np.ones(1, np.float32)
  n = np.arange(M)
  denom = (M - 1) if sym else M
  return (0.54 - 0.46 * np.cos(2.0 * np.pi * n / denom)).astype(np.float32)

blackman 

blackman(M: int, sym: bool = False) -> np.ndarray

Blackman window. sym=False = periodic (STFT) form.

Source code in aneforge/dsp.py 
37
38
39
40
41
42
43
44
def blackman(M: int, sym: bool = False) -> np.ndarray:
  """Blackman window. `sym=False` = periodic (STFT) form."""
  if M == 1: return np.ones(1, np.float32)
  n = np.arange(M)
  denom = (M - 1) if sym else M
  a0, a1, a2 = 0.42, 0.5, 0.08
  return (a0 - a1 * np.cos(2.0 * np.pi * n / denom)
          + a2 * np.cos(4.0 * np.pi * n / denom)).astype(np.float32)

kaiser 

kaiser(M: int, beta: float = 8.6, sym: bool = False) -> np.ndarray

Kaiser window (I0 Bessel ratio, np.i0). Default beta=8.6 is Blackman-like. sym=False = periodic (STFT) form.

Source code in aneforge/dsp.py 
47
48
49
50
51
52
def kaiser(M: int, beta: float = 8.6, sym: bool = False) -> np.ndarray:
  """Kaiser window (I0 Bessel ratio, np.i0). Default beta=8.6 is Blackman-like. `sym=False` = periodic (STFT) form."""
  if M == 1: return np.ones(1, np.float32)
  n = np.arange(M)
  alpha = ((M - 1) if sym else M) / 2.0
  return (np.i0(beta * np.sqrt(1.0 - ((n - alpha) / alpha) ** 2)) / np.i0(beta)).astype(np.float32)

bartlett 

bartlett(M: int, sym: bool = False) -> np.ndarray

Bartlett (triangular, zero endpoints) window. sym=False = periodic (STFT) form.

Source code in aneforge/dsp.py 
55
56
57
58
59
60
def bartlett(M: int, sym: bool = False) -> np.ndarray:
  """Bartlett (triangular, zero endpoints) window. `sym=False` = periodic (STFT) form."""
  if M == 1: return np.ones(1, np.float32)
  n = np.arange(M)
  half = ((M - 1) if sym else M) / 2.0
  return (1.0 - np.abs(n - half) / half).astype(np.float32)

tukey 

tukey(M: int, alpha: float = 0.5, sym: bool = False) -> np.ndarray

Tukey (tapered cosine) window; alpha is the taper fraction (0 = boxcar, 1 = hann). sym=False = periodic form.

Source code in aneforge/dsp.py 
63
64
65
66
67
68
69
70
71
72
def tukey(M: int, alpha: float = 0.5, sym: bool = False) -> np.ndarray:
  """Tukey (tapered cosine) window; alpha is the taper fraction (0 = boxcar, 1 = hann). `sym=False` = periodic form."""
  if M == 1: return np.ones(1, np.float32)
  if alpha <= 0: return np.ones(M, np.float32)
  if alpha >= 1: return hann(M, sym=sym)
  N = M if sym else M + 1                                # periodic = symmetric on M+1 points, last dropped
  r = np.arange(N) / (N - 1)
  t = np.minimum(r, 1.0 - r)                             # distance to the nearer edge; taper is symmetric
  w = np.where(t < alpha / 2.0, 0.5 * (1.0 + np.cos(np.pi * (2.0 * t / alpha - 1.0))), 1.0)
  return w[:M].astype(np.float32)

get_window 

get_window(window, M: int) -> np.ndarray

Resolve window (name str, 'boxcar'/None for rectangular, or an array) to a length-M fp32 coefficient vector.

Source code in aneforge/dsp.py 
79
80
81
82
83
84
85
86
87
88
89
90
def get_window(window, M: int) -> np.ndarray:
  """Resolve `window` (name str, 'boxcar'/None for rectangular, or an array) to a length-M fp32 coefficient vector."""
  if window is None or (isinstance(window, str) and window in ("boxcar", "rect")):
    return np.ones(M, np.float32)                        # `in` on the tuple would choke on array inputs
  if isinstance(window, str):
    if window not in _WINDOWS:
      raise ValueError(f"unknown window {window!r}; choose from {list(_WINDOWS)} or 'boxcar'")
    return _WINDOWS[window](M, sym=False)
  w = np.asarray(window, np.float32)
  if w.shape != (M,):
    raise ValueError(f"window array must have length {M}; got {w.shape}")
  return w

sawtooth 

sawtooth(t, width: float = 1.0) -> np.ndarray

Periodic sawtooth/triangle waveform, period 2*pi (scipy.signal.sawtooth). width in [0,1] sets the rising fraction (1 = rising ramp, 0.5 = triangle, 0 = falling).

Source code in aneforge/dsp.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def sawtooth(t, width: float = 1.0) -> np.ndarray:
  """Periodic sawtooth/triangle waveform, period 2*pi (scipy.signal.sawtooth). `width` in [0,1] sets the rising fraction (1 = rising ramp, 0.5 = triangle, 0 = falling)."""
  t = np.asarray(t, np.float64)
  w = np.broadcast_to(np.asarray(width, np.float64), t.shape)
  tmod = np.mod(t, 2.0 * np.pi)
  valid = (w >= 0.0) & (w <= 1.0)
  rise = valid & (tmod < w * 2.0 * np.pi)
  y = np.full(t.shape, np.nan, np.float64)
  with np.errstate(divide="ignore", invalid="ignore"):
    np.copyto(y, tmod / (np.pi * w) - 1.0, where=rise)
    np.copyto(y, (np.pi * (w + 1.0) - tmod) / (np.pi * (1.0 - w)), where=valid & ~rise)
  return y.astype(np.float32)

square 

square(t, duty: float = 0.5) -> np.ndarray

Periodic square wave, period 2*pi (scipy.signal.square): +1 for the first duty fraction of the cycle, -1 after. duty in [0,1].

Source code in aneforge/dsp.py 
109
110
111
112
113
114
115
116
117
118
119
def square(t, duty: float = 0.5) -> np.ndarray:
  """Periodic square wave, period 2*pi (scipy.signal.square): +1 for the first `duty` fraction of the cycle, -1 after. `duty` in [0,1]."""
  t = np.asarray(t, np.float64)
  w = np.broadcast_to(np.asarray(duty, np.float64), t.shape)
  tmod = np.mod(t, 2.0 * np.pi)
  valid = (w >= 0.0) & (w <= 1.0)
  high = valid & (tmod < w * 2.0 * np.pi)
  y = np.full(t.shape, np.nan, np.float64)
  np.copyto(y, 1.0, where=high)
  np.copyto(y, -1.0, where=valid & ~high)
  return y.astype(np.float32)

chirp 

chirp(t, f0: float, t1: float, f1: float, method: str = 'linear') -> np.ndarray

Frequency-swept cosine (scipy.signal.chirp, linear sweep): frequency ramps from f0 at t=0 to f1 at t=t1. Phase = 2pi(f0t + (f1-f0)t^2/(2*t1)).

Source code in aneforge/dsp.py 
122
123
124
125
126
127
128
129
def chirp(t, f0: float, t1: float, f1: float, method: str = "linear") -> np.ndarray:
  """Frequency-swept cosine (scipy.signal.chirp, linear sweep): frequency ramps from f0 at t=0 to f1 at t=t1. Phase = 2*pi*(f0*t + (f1-f0)*t^2/(2*t1))."""
  if method not in ("linear", "lin", "li"):
    raise ValueError(f"chirp: only method='linear' is implemented; got {method!r}")
  t = np.asarray(t, np.float64)
  beta = (float(f1) - float(f0)) / float(t1)
  phase = 2.0 * np.pi * (f0 * t + 0.5 * beta * t * t)
  return np.cos(phase).astype(np.float32)

fir_filter 

fir_filter(x, taps, mode: str = 'same')

FIR filter y = x * taps (true convolution) on the ANE; taps flipped (ANE conv is correlation). mode in full/same/valid/lfilter.

Source code in aneforge/dsp.py 
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
def fir_filter(x, taps, mode: str = "same"):
  """FIR filter y = x * taps (true convolution) on the ANE; taps flipped (ANE conv is correlation). `mode` in full/same/valid/lfilter."""
  x = np.asarray(x, np.float32).ravel()
  taps = np.asarray(taps, np.float32).ravel()
  L, K = x.shape[0], taps.shape[0]
  if K > L:
    raise ValueError(f"fir_filter: {K} taps longer than signal length {L}")

  if K > _MAX_CONV_TAPS:                              # long kernels -> FFT convolution
    full = fft_convolve(x, taps)                     # length L+K-1
    return _trim_conv(full, L, K, mode)

  # short kernels: native conv (== correlation), so flip taps and zero-pad to the requested slice.
  hflip = np.ascontiguousarray(taps[::-1]).astype(np.float16).reshape(1, 1, 1, K)
  if mode in ("full", "lfilter"): left = K - 1
  elif mode == "same": left = (K - 1) // 2
  elif mode == "valid": left = 0
  else: raise ValueError(f"fir_filter: mode must be full/same/valid/lfilter; got {mode!r}")
  if mode == "full": right = K - 1
  elif mode == "same": right = K - 1 - left
  elif mode == "lfilter": right = 0
  else:                                                # valid
    right = 0
  xp = np.concatenate([np.zeros(left, np.float32), x, np.zeros(right, np.float32)])
  Lp = xp.shape[0]
  inp = af.input((1, 1, 1, Lp))
  model = af.compile(af.conv(inp, hflip, pad=0))
  y = model(xp.astype(np.float16).reshape(1, 1, 1, Lp))
  return y.ravel().astype(np.float32)

fft_convolve 

fft_convolve(x, h, block: int | None = None)

Linear convolution y = x * h via the FFT, length L+K-1; long signals use overlap-add against the cached kernel spectrum.

Source code in aneforge/dsp.py 
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
def fft_convolve(x, h, block: int | None = None):
  """Linear convolution y = x * h via the FFT, length L+K-1; long signals use overlap-add against the cached kernel spectrum."""
  x = np.asarray(x, np.float32).ravel()
  h = np.asarray(h, np.float32).ravel()
  L, K = x.shape[0], h.shape[0]
  out_len = L + K - 1

  # single-block path: one FFT covers the whole linear convolution.
  if block is None:
    single = _next_fft_size(out_len)
    if L <= 4 * K or single <= 1024:
      return _fft_conv_block(x, h, single)[:out_len].astype(np.float32)
    block = single  # (not reached for the common case; kept for explicitness)

  # overlap-add: choose an FFT size, step = N-K+1 samples per block.
  N = _next_fft_size(max(block, 2 * K))
  step = N - K + 1
  Hr, Hi = fft(np.pad(h, (0, N - K)), np.zeros(N, np.float32), N)
  fplan = fft_plan(N); iplan = ifft_plan(N)
  y = np.zeros(out_len, np.float32)
  for start in range(0, L, step):
    seg = x[start:start + step]
    seg = np.pad(seg, (0, N - seg.shape[0]))
    Xr, Xi = fplan(seg, np.zeros(N, np.float32))
    Yr = Xr * Hr - Xi * Hi
    Yi = Xr * Hi + Xi * Hr
    yr = _ifft_real_scaled(Yr, Yi, N, iplan)
    end = min(start + N, out_len)
    y[start:end] += yr[:end - start].astype(np.float32)
  return y

freq_filter 

freq_filter(x, kind: str, cutoff, fs: float = 2.0)

Brick-wall frequency-domain filter: FFT, zero the rejected bins, inverse-FFT. kind in lowpass/highpass/bandpass/bandstop.

Source code in aneforge/dsp.py 
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def freq_filter(x, kind: str, cutoff, fs: float = 2.0):
  """Brick-wall frequency-domain filter: FFT, zero the rejected bins, inverse-FFT. `kind` in lowpass/highpass/bandpass/bandstop."""
  x = np.asarray(x, np.float32).ravel()
  L = x.shape[0]
  N = _next_fft_size(L)
  freqs = np.fft.fftfreq(N, d=1.0 / fs)                # bin center frequencies
  absf = np.abs(freqs)

  if kind in ("lowpass", "highpass"):
    fc = float(np.asarray(cutoff).ravel()[0])
    passband = (absf <= fc) if kind == "lowpass" else (absf >= fc)
  elif kind in ("bandpass", "bandstop"):
    lo, hi = (float(c) for c in np.asarray(cutoff).ravel()[:2])
    inband = (absf >= lo) & (absf <= hi)
    passband = inband if kind == "bandpass" else ~inband
  else:
    raise ValueError(f"freq_filter: kind must be lowpass/highpass/bandpass/bandstop; got {kind!r}")
  mask = passband.astype(np.float32)

  Xr, Xi = fft(np.pad(x, (0, N - L)), np.zeros(N, np.float32), N)
  Xr = Xr * mask
  Xi = Xi * mask
  yr = _ifft_real_scaled(Xr, Xi, N)                   # fp16-range-guarded iFFT
  return yr[:L].astype(np.float32)

stft 

stft(x, win=256, hop=None, window: str = 'hann')

Short-time Fourier transform: window each frame, FFT, stack. Returns (Zr, Zi), each [n_freq, n_frames] (no 1/sum(win) scaling).

Source code in aneforge/dsp.py 
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
def stft(x, win=256, hop=None, window: str = "hann"):
  """Short-time Fourier transform: window each frame, FFT, stack. Returns (Zr, Zi), each [n_freq, n_frames] (no 1/sum(win) scaling)."""
  x = np.asarray(x, np.float32).ravel()
  if isinstance(win, (int, np.integer)):
    win_len = int(win)
    w = get_window(window, win_len)
  else:
    w = np.asarray(win, np.float32).ravel()
    win_len = w.shape[0]
  if hop is None: hop = win_len // 4
  N = _next_fft_size(win_len)
  n_freq = win_len // 2 + 1

  frames = _frame(x, win_len, hop) * w                 # [n_frames, win_len], windowed
  n_frames = frames.shape[0]
  plan = fft_plan(N)
  Zr = np.empty((n_freq, n_frames), np.float32)
  Zi = np.empty((n_freq, n_frames), np.float32)
  zero = np.zeros(N, np.float32)
  for t in range(n_frames):
    fr = np.pad(frames[t], (0, N - win_len))
    Fr, Fi = plan(fr, zero)
    Zr[:, t] = Fr[:n_freq]
    Zi[:, t] = Fi[:n_freq]
  return Zr, Zi

spectrogram 

spectrogram(x, win=256, hop=None, window: str = 'hann', mode: str = 'magnitude')

Magnitude (or power) spectrogram = |STFT|; mode in {'magnitude','power'}.

Source code in aneforge/dsp.py 
304
305
306
307
308
def spectrogram(x, win=256, hop=None, window: str = "hann", mode: str = "magnitude"):
  """Magnitude (or power) spectrogram = |STFT|; `mode` in {'magnitude','power'}."""
  Zr, Zi = stft(x, win=win, hop=hop, window=window)
  p = Zr.astype(np.float32) ** 2 + Zi.astype(np.float32) ** 2
  return p if mode == "power" else np.sqrt(p)

correlate 

correlate(a, b, mode: str = 'valid')

Cross-correlation r[k] = sum_n a[n+k] * b[n] on the ANE; short templates (Lb<=15) use the CrossCorrelation bridge, wider ones FFT. mode in valid/full/same.

Source code in aneforge/dsp.py 
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
def correlate(a, b, mode: str = "valid"):
  """Cross-correlation r[k] = sum_n a[n+k] * b[n] on the ANE; short templates (Lb<=15) use the CrossCorrelation bridge, wider ones FFT. `mode` in valid/full/same."""
  a = np.asarray(a, np.float32).ravel()
  b = np.asarray(b, np.float32).ravel()
  La, Lb = a.shape[0], b.shape[0]
  if Lb >= La:
    raise ValueError(f"correlate: template length {Lb} must be < signal length {La} "
                     f"(CrossCorrelation bridge requires a strictly smaller template)")
  if mode == "valid": ap = a
  elif mode == "full":
    ap = np.concatenate([np.zeros(Lb - 1, np.float32), a, np.zeros(Lb - 1, np.float32)])
  elif mode == "same":
    pad = (Lb - 1) // 2
    ap = np.concatenate([np.zeros(pad, np.float32), a, np.zeros(Lb - 1 - pad, np.float32)])
  else:
    raise ValueError(f"correlate: mode must be valid/full/same; got {mode!r}")

  # wide template -> FFT: correlate(ap,b) == full-convolve(ap, reverse(b)) sliced.
  if Lb > _MAX_CONV_TAPS:
    full = fft_convolve(ap, b[::-1])                 # length len(ap)+Lb-1
    return full[Lb - 1:ap.shape[0]].astype(np.float32)

  Wp = ap.shape[0]
  model = af.compile(af.cross_correlation(af.input((1, Wp)), af.input((1, Lb))))
  y = model(ap.astype(np.float16).reshape(1, Wp), b.astype(np.float16).reshape(1, Lb))
  return y.ravel().astype(np.float32)

autocorrelate 

autocorrelate(x, max_lag: int | None = None)

Autocorrelation r[k] for lags k=0..max_lag, on the ANE (correlation against the signal's own prefix); max_lag defaults to len(x)//4.

Source code in aneforge/dsp.py 
341
342
343
344
345
346
347
348
349
350
def autocorrelate(x, max_lag: int | None = None):
  """Autocorrelation r[k] for lags k=0..max_lag, on the ANE (correlation against the signal's own prefix); `max_lag` defaults to len(x)//4."""
  x = np.asarray(x, np.float32).ravel()
  L = x.shape[0]
  if max_lag is None: max_lag = L // 4
  tmpl_len = L - max_lag
  if tmpl_len < 1 or tmpl_len >= L:
    raise ValueError(f"autocorrelate: max_lag={max_lag} out of range for length {L}")
  tmpl = x[:tmpl_len]
  return correlate(x, tmpl, mode="valid")

iir_filter 

iir_filter(x, b, a, n_taps: int = 256)

Arch-limited IIR via a fixed-length FIR unroll: truncate the IIR impulse response to n_taps and run as FIR (a direct IIR is the scan recurrence the ANE lacks).

Source code in aneforge/dsp.py 
355
356
357
358
359
360
361
362
363
364
def iir_filter(x, b, a, n_taps: int = 256):
  """Arch-limited IIR via a fixed-length FIR unroll: truncate the IIR impulse response to `n_taps` and run as FIR (a direct IIR is the scan recurrence the ANE lacks)."""
  import scipy.signal as ss
  b = np.asarray(b, np.float64).ravel()
  a = np.asarray(a, np.float64).ravel()
  # truncated impulse response (response to a unit impulse, n_taps long)
  imp = np.zeros(n_taps, np.float64)
  imp[0] = 1.0
  ir = ss.lfilter(b, a, imp).astype(np.float32)        # type: ignore[union-attr]  # lfilter returns ndarray; stubs over-constrain return
  return fir_filter(x, ir, mode="lfilter")

hilbert 

hilbert(x)

Analytic signal via the on-ANE FFT: z(t) = x(t) + i * Hx.

Returns (z_re, z_im) as numpy arrays: z_re == x and z_im is the Hilbert transform. Any length >= 2 works -- the staged FFT plans for exactly N and does not pad (a prime N degenerates to one dense N x N DFT, as everywhere else in aneforge.fft). Oracle: scipy.signal.hilbert.

Source code in aneforge/dsp.py 
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
def hilbert(x):
  """Analytic signal via the on-ANE FFT: z(t) = x(t) + i * H[x](t).

  Returns (z_re, z_im) as numpy arrays: z_re == x and z_im is the Hilbert transform.
  Any length >= 2 works -- the staged FFT plans for exactly N and does not pad (a prime N
  degenerates to one dense N x N DFT, as everywhere else in aneforge.fft).
  Oracle: `scipy.signal.hilbert`.
  """
  x = np.asarray(x, np.float32).ravel()
  N = x.shape[0]
  if N < 2: raise ValueError(f"hilbert: need at least 2 samples; got {N}")
  # one-sided response in the freq domain. Even N has a Nyquist bin to halve, odd N does not --
  # applying the even branch to an odd length double-counts the top bin (relerr 2.6e-02 vs 1.2e-03
  # at N=255). Same split as scipy.signal.hilbert.
  h = np.zeros(N, np.float32)
  h[0] = 1.0
  if N % 2 == 0:
    h[1:N // 2] = 2.0
    h[N // 2] = 1.0
  else:
    h[1:(N + 1) // 2] = 2.0
  # FFT of x, apply the (real) one-sided response, IFFT back to the analytic signal
  Xr, Xi = fft(x, np.zeros(N, np.float32), N)
  y_re, y_im = ifft(Xr * h, Xi * h, N)
  return y_re.astype(np.float32), y_im.astype(np.float32)

welch 

welch(x, fs: float = 1.0, nperseg: int = 256, noverlap=None, window: str = 'hann', scaling: str = 'density')

Welch power spectral density from the on-ANE stft: window, segment, average the squared magnitudes, normalize. Matches scipy.signal.welch(..., detrend=False) (no per-segment detrending). scaling: 'density' (1/(fs*sum(w^2))) or 'spectrum' (1/sum(w)^2); the one-sided spectrum doubles interior bins. Returns (f, Pxx).

Source code in aneforge/dsp.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
def welch(x, fs: float = 1.0, nperseg: int = 256, noverlap=None, window: str = "hann",
          scaling: str = "density"):
  """Welch power spectral density from the on-ANE `stft`: window, segment, average the squared
  magnitudes, normalize. Matches `scipy.signal.welch(..., detrend=False)` (no per-segment
  detrending). `scaling`: 'density' (1/(fs*sum(w^2))) or 'spectrum' (1/sum(w)^2); the one-sided
  spectrum doubles interior bins. Returns (f, Pxx)."""
  if nperseg & (nperseg - 1):
    raise ValueError("welch: nperseg must be a power of two (the staged on-ANE FFT pads to one, which would shift the bins)")
  if noverlap is None: noverlap = nperseg // 2
  w = get_window(window, nperseg)
  Zr, Zi = stft(x, win=nperseg, hop=nperseg - noverlap, window=window)   # [n_freq, n_frames], FFTs on the engine
  P = (Zr.astype(np.float64) ** 2 + Zi.astype(np.float64) ** 2).mean(axis=1)
  if scaling == "density":
    P /= fs * float(np.sum(w.astype(np.float64) ** 2))
  elif scaling == "spectrum":
    P /= float(np.sum(w.astype(np.float64))) ** 2
  else:
    raise ValueError(f"welch: scaling={scaling!r} (use 'density' or 'spectrum')")
  P[1:-1 if nperseg % 2 == 0 else None] *= 2.0             # one-sided: double all but DC (and Nyquist if present)
  f = np.arange(P.shape[0], dtype=np.float64) * (fs / nperseg)
  return f, P.astype(np.float32)