Skip to content

Symbolic arithmetic

One signature, no annotations in the body. Every comment below is the shape numty infers, as a formula over the names in the signature:

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view
from numty import F64, I64

def features(
    signal: F64["batch time"],
    image: F64["batch chan height width"],
    offsets: I64["batch+1"],
) -> None:
    np.diff(signal, axis=1)                     # (batch, time-1)
    np.fft.rfft(signal, axis=1)                 # (batch, time//2+1), complex128
    sliding_window_view(signal, 5, axis=1)      # (batch, time-4, 5)
    np.pad(signal, ((0, 0), (2, 2)))            # (batch, time+4)
    np.concatenate([signal, signal], axis=1)    # (batch, 2*time)
    image.reshape(image.shape[0], -1)           # (batch, chan*height*width)
    np.diff(offsets)                            # (batch,)
    signal[:, 1:] - signal[:, :-1]              # (batch, time-1)
    signal[:, 1:] - signal                      # off by one
error: broadcast shapes disagree at trailing axis 0: time-1 vs time (shapes (batch, time-1) vs (batch, time))  [numty-broadcast]

batch, time, chan, height and width are never declared: the strings mint them. The sizes are unknown until the function is called, and the off-by-one on the last line is still an error, because time-1 cannot equal time whatever time is.

What is tracked

Dims stay exact through arithmetic: repeat/tile multiply (N*2), pad adds (N+2), convolve adds lengths (N+M-1), rfft halves (N//2+1), split divides (N//2), and reshape proves N*2 vs 2*N (dividing -1 exactly); mismatches error with readable formulas. Products of distinct symbolic dims track as monomials (H*D, C*H*W): flattening (B, C, H, W) with -1 infers (B, C*H*W), (N, K) reshaped to (N, -1) infers (N, K), and the round-trip back proves clean by commutativity. A reshape target may spell its sizes as arithmetic over known ones: x.reshape(x.shape[0] * x.shape[1], -1) on (N, M, K) is (M*N, K), and 2 * n, n + 1 and n // 2 are read the same way. sliding_window_view appends the window after the shrunken dims (N-w+1, symbolic included — the conv-output formula falls out). Linear forms also spell in annotations: I64["N+1"] denotes one more than N, so the offsets/lengths duality checks both ways — np.diff of I64["N+1"] is I64[N], and np.concatenate([[0], np.cumsum(lengths)]) of I64[N] is I64["N+1"] (no type: ignore). Supported per dim: N+1, N-1, 2*N, N+M, plain constants ("5" is Literal[5]); anything non-linear (N*M, N//2) fails loudly instead of going quietly Any.

What that unlocks

Code that failed before, exact now:

def f(u: F64[N], w: F64[N, K]) -> None:
    n: N = u.shape[0]
    # Was: "reshape changes total size: 1 -> 2". Now: exact F64[2, N].
    r: F64[2, N] = np.reshape(np.repeat(u, 2), (2, n))
    # -1 divides the canonical product (N*2 / 2 = N); the quotient
    # contracts exactly downstream. Was: bailed to Any.
    t = np.reshape(np.repeat(u, 2), (2, -1))
    m: F64[2, K] = np.dot(t, w)
    # rfft on symbolic input claimed (N,) — silently wrong.
    # Now: (N//2+1), and it participates in checks.
    s = np.fft.rfft(u)