pipekit-array

Array-API operators on top of pipekit. One operator runs on numpy, JAX, CuPy, PyTorch, or dask via array_namespace(x) dispatch — see the design doc for the full architecture, ADRs, and the v0.1 operator catalogue.

Phase A of the v0.1 implementation has landed: namespace dispatch + reductions + stack/concat combinators. Phases B–D add geometry, QC, and the trained-model wrapper ModelOp.

array_namespace

pipekit_array._namespace

Array API namespace dispatch shim.

The single point of backend dispatch for every pipekit-array operator. array_namespace(*xs) returns the Array API namespace (numpy, jax.numpy, torch via array-api-compat, …) shared by all input arrays.

Tries the optional array-api-compat library first — without it, PyTorch tensors don't expose __array_namespace__ natively and torch backends won't dispatch. Falls back to native x.__array_namespace__() lookup if the compat shim isn't installed.

See docs/design/architecture.md §4 and ADR A1 for the rationale.

array_namespace(*xs)

Return the Array API namespace shared by all input arrays.

Parameters:

Name	Type	Description	Default
*xs	Any	One or more array-like objects. All must share the same Array API backend.	()

Returns:

Type	Description
Any	The module-like namespace object (e.g. numpy, jax.numpy,
Any	array_api_compat.torch).

Raises:

Type	Description
TypeError	If no input implements the Array API, or if the inputs span more than one backend. The error message is consistent whether or not array-api-compat is installed.

Source code in packages/pipekit-array/src/pipekit_array/_namespace.py

def array_namespace(*xs: Any) -> Any:
    """Return the Array API namespace shared by all input arrays.

    Args:
        *xs: One or more array-like objects. All must share the same
            Array API backend.

    Returns:
        The module-like namespace object (e.g. `numpy`, `jax.numpy`,
        `array_api_compat.torch`).

    Raises:
        TypeError: If no input implements the Array API, or if the
            inputs span more than one backend. The error message is
            consistent whether or not `array-api-compat` is installed.
    """
    compat_ns: Callable[..., Any] | None = None
    try:
        import array_api_compat  # ty: ignore[unresolved-import]

        compat_ns = array_api_compat.array_namespace
    except ImportError:
        pass

    if compat_ns is not None:
        try:
            return compat_ns(*xs)
        except TypeError as exc:
            # array-api-compat's error messages differ from our
            # contract; re-raise with stable wording so tests and
            # users see a consistent shape.
            msg = str(exc)
            if "Multiple namespaces" in msg:
                names = sorted({type(x).__module__.split(".")[0] for x in xs})
                raise TypeError(_multi_backend_msg(names)) from exc
            raise TypeError(_NO_BACKEND_MSG) from exc

    namespaces: set[Any] = set()
    for x in xs:
        ns = getattr(x, "__array_namespace__", None)
        if ns is None:
            raise TypeError(
                f"Input of type {type(x).__name__!r} does not implement "
                f"the Array API (`__array_namespace__` missing). " + _NO_BACKEND_MSG
            )
        namespaces.add(ns())
    if len(namespaces) == 0:
        raise TypeError(_NO_BACKEND_MSG)
    if len(namespaces) > 1:
        names = sorted({ns.__name__ for ns in namespaces})
        raise TypeError(_multi_backend_msg(names))
    return next(iter(namespaces))

Reductions

pipekit_array.reduce

Reductions — MeanScalar.

Wraps xp.mean as a pipekit Operator. See docs/design/api/operators.md for the spec.

MeanScalar

Bases: Operator

Reduce an array to a scalar (or per-axis means) via xp.mean.

Defaults to a full reduction (axis=None) yielding a scalar of rank 0. Output dtype matches the input's floating dtype per the Array API mean contract.

Parameters:

Name	Type	Description	Default
axis	int | tuple[int, ...] | None	Axis or tuple of axes to reduce over. None reduces all axes.	None
keepdims	bool	If True, retains the reduced axes with size 1.	False

Raises:

Type	Description
ValueError	From compute_output_signature when axis is out of range for the input signature's rank (xp.mean would raise the backend-native equivalent at apply time).

Example::

out = MeanScalar()(image)               # scalar
out = MeanScalar(axis=(-2, -1))(image)  # per-band mean
out = MeanScalar(axis=0, keepdims=True)(image)  # shape (1, H, W)

Source code in packages/pipekit-array/src/pipekit_array/reduce.py

class MeanScalar(Operator):
    """Reduce an array to a scalar (or per-axis means) via `xp.mean`.

    Defaults to a full reduction (`axis=None`) yielding a scalar of
    rank 0. Output dtype matches the input's floating dtype per the
    Array API `mean` contract.

    Args:
        axis: Axis or tuple of axes to reduce over. `None` reduces
            all axes.
        keepdims: If True, retains the reduced axes with size 1.

    Raises:
        ValueError: From `compute_output_signature` when ``axis`` is out
            of range for the input signature's rank (``xp.mean`` would
            raise the backend-native equivalent at apply time).

    Example::

        out = MeanScalar()(image)               # scalar
        out = MeanScalar(axis=(-2, -1))(image)  # per-band mean
        out = MeanScalar(axis=0, keepdims=True)(image)  # shape (1, H, W)
    """

    def __init__(
        self,
        axis: int | tuple[int, ...] | None = None,
        keepdims: bool = False,
    ) -> None:
        # a tuple of axes reloads from JSON as a list, which `xp.mean` rejects
        self.axis = tuple(axis) if isinstance(axis, list) else axis
        self.keepdims = keepdims

    def _apply(self, x: Any) -> Any:
        xp = array_namespace(x)
        return xp.mean(x, axis=self.axis, keepdims=self.keepdims)

    def get_config(self) -> dict[str, Any]:
        return {"axis": self.axis, "keepdims": self.keepdims}

    def compute_output_signature(self, input_signature: Any) -> Any:
        if input_signature is None:
            return None

        input_sig: Signature = input_signature
        n_dims = len(input_sig.dims)

        def _validate(axis: int) -> int:
            """Return positive axis index; raise on out-of-range."""
            if n_dims == 0 or axis < -n_dims or axis >= n_dims:
                raise ValueError(
                    f"MeanScalar.axis={axis!r} is out of range for input "
                    f"signature with {n_dims} dim(s); xp.mean would raise."
                )
            return axis % n_dims

        if self.axis is None:
            axes_to_reduce: tuple[int, ...] = tuple(range(n_dims))
        elif isinstance(self.axis, int):
            axes_to_reduce = (_validate(self.axis),)
        else:
            axes_to_reduce = tuple(_validate(a) for a in self.axis)

        if self.keepdims:
            new_dims = tuple(
                (name, 1) if i in axes_to_reduce else (name, size)
                for i, (name, size) in enumerate(input_sig.dims)
            )
        else:
            new_dims = tuple(
                (name, size)
                for i, (name, size) in enumerate(input_sig.dims)
                if i not in axes_to_reduce
            )
        return Signature(dims=new_dims, dtype=input_sig.dtype)

Combinators

pipekit_array.combinators

Combinators — StackAlong, ConcatenateAlong.

Thin wrappers over xp.stack and xp.concat (the Array API standardised names) so the operations compose into a pipeline. See docs/design/api/operators.md for the spec.

ApplyToBands lands in Phase B alongside BatchedMap.

ConcatenateAlong

Bases: Operator

Concatenate a list-of-arrays along an existing axis via xp.concat.

xp.concat is the Array API standardised name; numpy 2.0+ exposes it as concat alongside the legacy concatenate alias. All input arrays must share shape (except along axis), dtype, and backend.

Parameters:

Name	Type	Description	Default
axis	int	Existing axis to concatenate along. Default 0.	0

Raises:

Type	Description
ValueError	From _apply when the input sequence is empty (also via array_namespace when the inputs mix backends).

Example::

out = ConcatenateAlong(axis=0)([arr_a, arr_b])  # shape (2*N, H, W)

Source code in packages/pipekit-array/src/pipekit_array/combinators.py

class ConcatenateAlong(Operator):
    """Concatenate a list-of-arrays along an existing axis via `xp.concat`.

    `xp.concat` is the Array API standardised name; numpy 2.0+
    exposes it as `concat` alongside the legacy `concatenate` alias.
    All input arrays must share shape (except along `axis`), dtype,
    and backend.

    Args:
        axis: Existing axis to concatenate along. Default 0.

    Raises:
        ValueError: From ``_apply`` when the input sequence is empty
            (also via `array_namespace` when the inputs mix backends).

    Example::

        out = ConcatenateAlong(axis=0)([arr_a, arr_b])  # shape (2*N, H, W)
    """

    def __init__(self, axis: int = 0) -> None:
        self.axis = axis

    def _apply(self, xs: Sequence[Any]) -> Any:
        if len(xs) == 0:
            raise ValueError("ConcatenateAlong requires at least one input array.")
        xp = array_namespace(*xs)
        return xp.concat(list(xs), axis=self.axis)

    def get_config(self) -> dict[str, Any]:
        return {"axis": self.axis}

StackAlong

Bases: Operator

Stack a list-of-arrays along a new axis via xp.stack.

All input arrays must share the same shape, dtype, and backend.

Parameters:

Name	Type	Description	Default
axis	int	Position of the new axis in the output. Default 0.	0

Raises:

Type	Description
ValueError	From _apply when the input sequence is empty (also via array_namespace when the inputs mix backends).

Example::

out = StackAlong(axis=0)([arr_a, arr_b, arr_c])  # shape (3, H, W)

Source code in packages/pipekit-array/src/pipekit_array/combinators.py

class StackAlong(Operator):
    """Stack a list-of-arrays along a new axis via `xp.stack`.

    All input arrays must share the same shape, dtype, and backend.

    Args:
        axis: Position of the new axis in the output. Default 0.

    Raises:
        ValueError: From ``_apply`` when the input sequence is empty
            (also via `array_namespace` when the inputs mix backends).

    Example::

        out = StackAlong(axis=0)([arr_a, arr_b, arr_c])  # shape (3, H, W)
    """

    def __init__(self, axis: int = 0) -> None:
        self.axis = axis

    def _apply(self, xs: Sequence[Any]) -> Any:
        if len(xs) == 0:
            raise ValueError("StackAlong requires at least one input array.")
        xp = array_namespace(*xs)
        return xp.stack(list(xs), axis=self.axis)

    def get_config(self) -> dict[str, Any]:
        return {"axis": self.axis}