zorch.utils.binary_field

Binary-field ⟷ F_2 bit-vector representation for any binary tower dtype (binary_field_ghash, binary_field_t*).

An element of GF(2^W) is a W-dimensional vector over F_2; [unpack] exposes that coefficient vector and [pack] rebuilds the element — the GF(2^W) ≅ F_2^W isomorphism the ring-switch / tensor-algebra kernels ride on.

It is a shift/mask, NOT a bitcast. A hardware bitcast is byte-granular (the finest binary-field bitcast is GF(2^128) ↔ GF(2^8); GF(2^128) → F_2 fails the compatible-width check), so reaching the individual F_2 coefficients needs the explicit unpack. The uint32 storage limbs the shift/mask rides on stay private. uint32 is the finest limb, so it keeps the bit kernels valid for the widest set of tower dtypes — [field_bit_width] needs the width to be a whole number of limbs, so a uint64 limb would reject the 32-bit tower level (binary_field_t5) — and keeps the ring-switch {0, 1} × limb products narrow.

field_bit_width

field_bit_width(dtype: Any) -> int

W: the GF(2)-dimension of dtype (= its storage bits).

Source code in zorch/utils/binary_field.py

def field_bit_width(dtype: Any) -> int:
    """`W`: the GF(2)-dimension of `dtype` (= its storage bits)."""
    width = fnp.dtype(dtype).itemsize * 8
    if width % _LIMB_BITS != 0:
        raise ValueError(
            f"{fnp.dtype(dtype).name} is {width} bits; the bit kernels work over "
            f"uint{_LIMB_BITS} limbs and need a multiple of {_LIMB_BITS}"
        )
    return width

byte_select_xor_reduce

byte_select_xor_reduce(
    selectors: Array, values: Array
) -> Array

XOR weights selected by row-major packed bytes.

selectors is a uint8 matrix (n, B/8), where byte j stores selector bits 8*j .. 8*j+7 least-significant-bit first. values is a binary-field vector (B,); the result (n,) XORs values[b] wherever row bit b is set.

Where a kernel exists, a 256-entry XOR table per byte position feeds the gather, matching flock-core's UniSkipFoldTable: the kernel reads B/8 selector bytes per row and never expands the (n, B) bit matrix. Backends without one keep a compact source-level expression as the portable oracle.

Source code in zorch/utils/binary_field.py

@frx.jit
def byte_select_xor_reduce(selectors: Array, values: Array) -> Array:
    """XOR weights selected by row-major packed bytes.

    `selectors` is a `uint8` matrix `(n, B/8)`, where byte `j` stores selector
    bits `8*j .. 8*j+7` least-significant-bit first. `values` is a binary-field
    vector `(B,)`; the result `(n,)` XORs `values[b]` wherever row bit `b` is set.

    Where a kernel exists, a 256-entry XOR table per byte position feeds the
    gather, matching flock-core's `UniSkipFoldTable`: the kernel reads `B/8`
    selector bytes per row and never expands the `(n, B)` bit matrix. Backends
    without one keep a compact source-level expression as the portable oracle.
    """
    if selectors.ndim != 2 or values.ndim != 1:
        raise ValueError("packed-byte bit selection expects 2D selectors and 1D values")
    if fnp.dtype(selectors.dtype) != fnp.dtype(fnp.uint8):
        raise ValueError(
            f"packed-byte selectors must have dtype uint8, got {selectors.dtype}"
        )
    if selectors.shape[0] == 0 or selectors.shape[1] == 0:
        raise ValueError("packed-byte bit selection requires a non-empty matrix")
    bits = selectors.shape[1] * 8
    if values.shape != (bits,):
        raise ValueError(
            f"values must have shape ({bits},) for {selectors.shape[1]} selector "
            f"bytes, got {values.shape}"
        )
    values_dtype = fnp.dtype(values.dtype)
    if not values_dtype.name.startswith("binary_field"):
        raise ValueError("bit-select XOR reduction values must use a binary field")

    limbs = field_bit_width(values.dtype) // _LIMB_BITS
    values_l = _to_limbs(values)
    backend = frx.default_backend()
    if backend == _CUDA:
        if selectors.shape[1] <= 16:
            out_l = _bit_select_packed_bytes_pallas(selectors, values_l, limbs)
        else:
            out_l = _bit_select_wide_packed_bytes_pallas(selectors, values_l, limbs)
    elif backend == _METAL:
        # One kernel covers every row width, so Metal needs no wide/narrow
        # split: the gather is `S` table loads per row either way.
        out_l = _bit_select_packed_bytes_ffi(selectors, values_l, limbs)
    else:
        shifts = fnp.arange(8, dtype=fnp.uint8)
        unpacked = ((selectors[:, :, None] >> shifts) & 1).reshape(
            selectors.shape[0], bits
        )
        out_l = lax.reduce_xor(
            unpacked[:, :, None].astype(_LIMB) * values_l[None, :, :], (1,)
        )
    return _from_limbs(out_l, values.dtype)

bit_select_xor_reduce

bit_select_xor_reduce(
    selectors: Array,
    values: Array,
    *,
    reduce: BitSelectReduction
) -> Array

Select binary-field values with packed bits and XOR-reduce one axis.

selectors may be a one-dimensional vector of GF(2^W) elements. With reduce="elements", values has the same shape and the result has shape (W,): result bit-slice b XORs values[i] wherever bit b of selector i is set. With reduce="bits", values has shape (W,) and the result has the selectors' shape: each result i XORs values[b] wherever bit b of selector i is set.

reduce="elements" also batches: a selectors-major values of shape (n, N) reduces all N columns against the shared selectors in one pass — the result is (N, W), row k the reduction against values[:, k]. A batched ring-switch open uses this so the packed witness (the selectors) is read once for all N claims instead of once per claim.

For reduce="bits", selectors may instead be an explicit Boolean/integer 0/1 matrix of shape (n, B) with values.shape == (B,). This is the form used when a consumer already holds unpacked witness rows.

Where a kernel exists, the lowering streams directly into the limb output. It never materializes the broadcast (n, W, L) selection; its working set is one fixed-size register block. Backends without one use the compact source-level expression so the primitive remains portable and easy to validate.

Source code in zorch/utils/binary_field.py

@partial(frx.jit, static_argnames="reduce")
def bit_select_xor_reduce(
    selectors: Array,
    values: Array,
    *,
    reduce: BitSelectReduction,
) -> Array:
    """Select binary-field values with packed bits and XOR-reduce one axis.

    `selectors` may be a one-dimensional vector of GF(2^W) elements.
    With `reduce="elements"`, `values` has the same shape and the result has
    shape `(W,)`: result bit-slice `b` XORs `values[i]` wherever bit `b` of
    selector `i` is set. With `reduce="bits"`, `values` has shape `(W,)` and
    the result has the selectors' shape: each result `i` XORs `values[b]`
    wherever bit `b` of selector `i` is set.

    `reduce="elements"` also batches: a selectors-major `values` of shape
    `(n, N)` reduces all `N` columns against the shared selectors in one pass —
    the result is `(N, W)`, row `k` the reduction against `values[:, k]`. A
    batched ring-switch open uses this so the packed witness (the selectors) is
    read once for all `N` claims instead of once per claim.

    For `reduce="bits"`, selectors may instead be an explicit Boolean/integer
    0/1 matrix of shape `(n, B)` with `values.shape == (B,)`. This is the form
    used when a consumer already holds unpacked witness rows.

    Where a kernel exists, the lowering streams directly into the limb output.
    It never materializes the broadcast `(n, W, L)` selection; its working set
    is one fixed-size register block. Backends without one use the compact
    source-level expression so the primitive remains portable and easy to
    validate.
    """
    if selectors.ndim not in (1, 2):
        raise ValueError(
            "bit-select XOR reduction expects 1D packed or 2D unpacked selectors"
        )
    if values.ndim not in (1, 2):
        raise ValueError(
            "bit-select XOR reduction expects 1D values, or a 2D (n, N) stack "
            'for batched reduce="elements"'
        )
    if selectors.shape[0] == 0:
        raise ValueError("bit-select XOR reduction requires at least one selector")
    values_dtype = fnp.dtype(values.dtype)
    if not values_dtype.name.startswith("binary_field"):
        raise ValueError("bit-select XOR reduction values must use a binary field")

    if selectors.ndim == 2:
        if reduce != "bits":
            raise ValueError('unpacked selectors only support reduce="bits"')
        bits = selectors.shape[1]
        if bits == 0 or values.shape != (bits,):
            raise ValueError(
                f"values must have shape ({bits},) for unpacked selectors, got "
                f"{values.shape}"
            )
        limbs = field_bit_width(values.dtype) // _LIMB_BITS
        values_l = _to_limbs(values)
        # Metal is deliberately absent: its accelerated lowering gathers packed
        # selector bytes, and these selectors are an explicit 0/1 matrix whose
        # bit count need not be a multiple of 8. Packing it to reach the kernel
        # would cost a pass over the (n, B) matrix the kernel exists to avoid
        # materializing, so Metal keeps the portable expression here.
        if frx.default_backend() == _CUDA:
            out_l = _bit_select_unpacked_bits_pallas(selectors, values_l, limbs)
        else:
            selected = selectors.astype(_LIMB)
            out_l = lax.reduce_xor(selected[:, :, None] * values_l[None, :, :], (1,))
        return _from_limbs(out_l, values.dtype)

    if fnp.dtype(selectors.dtype) != values_dtype:
        raise ValueError(
            "bit-select XOR reduction requires selectors and values of the same "
            "binary-field dtype"
        )
    width = field_bit_width(selectors.dtype)
    limbs = width // _LIMB_BITS
    selectors_l = _to_limbs(selectors)
    values_l = _to_limbs(values)

    if reduce == "elements":
        # `values` is `(n,)` for one claim or a selectors-major `(n, N)` stack
        # for N claims that share the selectors (the ring-switch batched-open
        # path: the packed witness is read once for all N). The reduced axis —
        # the one matching the selectors — leads in both.
        batched = values.ndim == 2
        if values.shape[0] != selectors.shape[0]:
            raise ValueError(
                f"values must match selectors on the reduced (leading) axis for "
                f'reduce="elements": {values.shape} vs {selectors.shape}'
            )
        backend = frx.default_backend()
        # One lowering, not two: both backends register a handler for this
        # target and XLA selects by platform. The set stays explicit because an
        # unregistered platform raises rather than falling back, so it has to
        # name exactly who has a handler. Metal appears in no arm below, which
        # is why its batched stacks and non-128 widths take the portable
        # expression until one is measured to matter.
        if backend in (_CUDA, _METAL) and not batched and width == 128:
            out_l = _bit_select_reduce_elements_ffi(selectors_l, values_l)
        elif backend == _CUDA:
            elements_pallas = (
                _bit_select_reduce_elements_batched_pallas
                if batched
                else _bit_select_reduce_elements_pallas
            )
            out_l = elements_pallas(selectors_l, values_l, width, limbs)
        else:
            bits = _bits(selectors)  # (n, width)
            if batched:
                # (n, width) x (n, N, limbs) -> (N, width, limbs)
                out_l = lax.reduce_xor(
                    bits[:, None, :, None] * values_l[:, :, None, :], (0,)
                )
            else:
                out_l = lax.reduce_xor(bits[:, :, None] * values_l[:, None, :], (0,))
        return _from_limbs(out_l, values.dtype)

    if reduce == "bits":
        if values.shape != (width,):
            raise ValueError(
                f'values must have shape ({width},) for reduce="bits", got '
                f"{values.shape}"
            )
        backend = frx.default_backend()
        if backend in (_CUDA, _METAL):
            selector_bytes = lax.bitcast_convert_type(selectors_l, fnp.uint8).reshape(
                selectors.shape[0], width // 8
            )
            packed_bytes = (
                _bit_select_packed_bytes_pallas
                if backend == _CUDA
                else _bit_select_packed_bytes_ffi
            )
            out_l = packed_bytes(selector_bytes, values_l, limbs)
        else:
            bits = _bits(selectors)
            out_l = lax.reduce_xor(bits[:, :, None] * values_l[None, :, :], (1,))
        return _from_limbs(out_l, values.dtype)

    raise ValueError(
        f'unknown reduction axis {reduce!r}; expected "bits" or "elements"'
    )

unpack

unpack(x: Array) -> Array

(...,) GF(2^W) -> (..., W) F_2: the element's F_2-coefficient vector, coefficient r at index r. The GF(2^W) ≅ F_2^W iso, realized as a shift/mask over the storage limbs (a bitcast cannot reach sub-byte coefficients) — the same bits [_bits] returns, retyped from uint32 to binary_field_t0.

Source code in zorch/utils/binary_field.py

def unpack(x: Array) -> Array:
    """`(...,)` GF(2^W) -> `(..., W)` F_2: the element's F_2-coefficient vector,
    coefficient `r` at index `r`. The GF(2^W) ≅ F_2^W iso, realized as a shift/mask
    over the storage limbs (a bitcast cannot reach sub-byte coefficients) — the
    same bits [`_bits`] returns, retyped from `uint32` to `binary_field_t0`."""
    return fnp.where(_bits(x).astype(bool), _f2(1), _f2(0))

pack

pack(coeffs: Array, dtype: Any) -> Array

(..., W) F_2 -> (...,) GF(2^W). Inverse of [unpack]: repack each 32 coefficients into a uint32 limb (Σ_r coeff_r · 2^r), then reinterpret.

Source code in zorch/utils/binary_field.py

def pack(coeffs: Array, dtype: Any) -> Array:
    """`(..., W)` F_2 -> `(...,)` GF(2^W). Inverse of [`unpack`]: repack each 32
    coefficients into a `uint32` limb (`Σ_r coeff_r · 2^r`), then reinterpret."""
    field_bit_width(dtype)  # validate W is a whole number of limbs
    bit = (coeffs == _f2(1)).astype(_LIMB)
    return _from_limbs(_limbs_from_bits(bit), dtype)