zorch.poly.univariate

Univariate polynomial evaluation, in evaluation and coefficient form.

powers

powers(x: Array, n: int) -> Array

(1, x, x², …, x^{n-1}) ascending, length n (n static).

Built by log-doubling (powers[m:2m] = powers[:m]·xᵐ) so the traced graph is O(log n), not n unrolled multiplies — a linear power chain makes the fused kernel's operand count scale with n and past a few thousand entries overruns the GPU's kernel-parameter space (the same cliff eval_coeffs avoids). The monomial-basis evaluation vector: ⟨coeffs, powers(x, n)⟩ = Σ cᵢ xⁱ.

Source code in zorch/poly/univariate.py

def powers(x: Array, n: int) -> Array:
    """``(1, x, x², …, x^{n-1})`` ascending, length ``n`` (``n`` static).

    Built by log-doubling (``powers[m:2m] = powers[:m]·xᵐ``) so the traced graph
    is O(log n), not ``n`` unrolled multiplies — a linear power chain makes the
    fused kernel's operand count scale with ``n`` and past a few thousand entries
    overruns the GPU's kernel-parameter space (the same cliff ``eval_coeffs``
    avoids). The monomial-basis evaluation vector: ``⟨coeffs, powers(x, n)⟩ =
    Σ cᵢ xⁱ``."""
    if n < 1:
        raise ValueError(f"powers needs n >= 1, got {n}")
    out = fnp.ones((1,), dtype=x.dtype)
    step = x
    while out.shape[0] < n:
        out = fnp.concatenate([out, out * step])
        step = step * step
    return out[:n]

eval_univariate

eval_univariate(evals: Array, x: Array) -> Array

Evaluate a univariate given by its values on [0, 1, ..., len-1] at x, by Lagrange interpolation over that integer domain.

A composer over the jitted basis kernel, so itself un-jitted.

Source code in zorch/poly/univariate.py

def eval_univariate(evals: Array, x: Array) -> Array:
    """Evaluate a univariate given by its values on ``[0, 1, ..., len-1]`` at
    ``x``, by Lagrange interpolation over that integer domain.

    A composer over the jitted basis kernel, so itself un-jitted."""
    nodes = naturals(evals.shape[0], evals.dtype)
    return fnp.dot(evals, compute_lagrange_basis(x, nodes))

compute_lagrange_basis

compute_lagrange_basis(r: Array, domain: Array) -> Array

All Lagrange basis evaluations L_{D,k}(r) over domain: L_{D,k}(r) = prod_{j != k} (r - x_j) / (x_k - x_j).

Direct form, not barycentric — barycentric divides by (r - node), which an r landing on a node would zero.

Source code in zorch/poly/univariate.py

@frx.jit
def compute_lagrange_basis(r: Array, domain: Array) -> Array:
    """All Lagrange basis evaluations ``L_{D,k}(r)`` over ``domain``:
    ``L_{D,k}(r) = prod_{j != k} (r - x_j) / (x_k - x_j)``.

    Direct form, not barycentric — barycentric divides by ``(r - node)``,
    which an ``r`` landing on a node would zero."""
    one = fnp.ones((), r.dtype)
    mask = fnp.eye(domain.shape[0], dtype=bool)
    numerators = fnp.prod(fnp.where(mask, one, (r - domain)[None, :]), axis=1)
    return numerators / _lagrange_denominators(domain)

compute_inv_vandermonde

compute_inv_vandermonde(degree: int, dtype: Any) -> Array

Inverse Vandermonde over the natural domain {0..degree}: coeffs = M @ evals for evals[j] = p(j).

Built in the base field — the Lagrange basis for an integer domain lives in the prime field — so one matrix serves BF and EF callers; EF evaluations promote at multiply time.

Source code in zorch/poly/univariate.py

def compute_inv_vandermonde(degree: int, dtype: Any) -> Array:
    """Inverse Vandermonde over the natural domain ``{0..degree}``:
    ``coeffs = M @ evals`` for ``evals[j] = p(j)``.

    Built in the base field — the Lagrange basis for an integer domain lives
    in the prime field — so one matrix serves BF and EF callers; EF
    evaluations promote at multiply time."""
    base = base_field(dtype)
    n = degree + 1
    one = fnp.array(1, base)
    zero = fnp.array(0, base)
    domain = naturals(n, dtype)
    denoms = _lagrange_denominators(domain)
    # Column j = coefficients of L_j(x) = prod_{k != j} (x - k) / denom_j,
    # expanded by repeated (x - k) multiplication over the coefficient list.
    columns = []
    for j in range(n):
        num_coeffs = [one]
        for k in range(n):
            if k != j:
                neg_k = -fnp.array(k, base)
                expanded = [zero] * (len(num_coeffs) + 1)
                for i, c in enumerate(num_coeffs):
                    expanded[i] = expanded[i] + c * neg_k
                    expanded[i + 1] = expanded[i + 1] + c
                num_coeffs = expanded
        columns.append(fnp.stack(num_coeffs) / denoms[j])
    return fnp.stack(columns, axis=1)

eval_coeffs

eval_coeffs(
    coeffs: Array,
    point: Array,
    *,
    schedule: str | None = None
) -> Array

p(point) = sum_i coeffs[..., i] * point**i — the coefficient-form dual of eval_univariate.

schedule: None/"auto" picks the measured per-field optimum; "horner" forces the unrolled chain (O(n)-deep graph — small n only); "scan" forces the log-depth prefix product. Byte-identical either way (exact field arithmetic). Horner fuses through a producer's pending expression stack; the scan keeps large degrees off both cliffs — a sequential scan's per-coefficient host launches and an unrolled power chain's kernel-parameter-space overflow.

Source code in zorch/poly/univariate.py

@partial(frx.jit, static_argnames=("schedule",))
def eval_coeffs(coeffs: Array, point: Array, *, schedule: str | None = None) -> Array:
    """``p(point) = sum_i coeffs[..., i] * point**i`` — the coefficient-form
    dual of ``eval_univariate``.

    ``schedule``: ``None``/``"auto"`` picks the measured per-field optimum;
    ``"horner"`` forces the unrolled chain (O(n)-deep graph — small ``n`` only);
    ``"scan"`` forces the log-depth prefix product. Byte-identical either way
    (exact field arithmetic). Horner fuses through a producer's pending
    expression stack; the scan keeps large degrees off both cliffs — a
    sequential scan's per-coefficient host launches and an unrolled power
    chain's kernel-parameter-space overflow."""
    if schedule not in (None, "auto", "horner", "scan"):
        raise ValueError(
            f"schedule must be None, 'auto', 'horner', or 'scan', got {schedule!r}"
        )
    n = coeffs.shape[-1]
    use_horner = (
        n <= _default_horner_max(coeffs.dtype)
        if schedule in (None, "auto")
        else schedule == "horner"
    )
    if use_horner:
        folded = coeffs[..., -1]
        for m in range(n - 2, -1, -1):
            folded = folded * point + coeffs[..., m]
        return folded
    # ``point**i`` as the prefix product of ``[1, point, point, …]`` (n entries),
    # laid out on the last axis so it dots the coefficients' degree axis directly.
    seq = fnp.where(
        fnp.arange(n) == 0,
        fnp.ones_like(point)[..., None],
        point[..., None],
    )
    powers = frx.lax.associative_scan(lambda a, b: a * b, seq, axis=-1)
    return fnp.sum(coeffs * powers, axis=-1)