Backend API#

Backend objects implement the numerical operations used by spaces, operators, functionals, and solvers.

Backend families#

spacecore.backend.BackendOps

Public numerical contract for SpaceCore backends.

spacecore.backend.NumpyOps

BackendOps implementation for the NumPy/SciPy ecosystem.

spacecore.backend.JaxOps

BackendOps implementation for the JAX ecosystem.

spacecore.backend.TorchOps

BackendOps implementation for PyTorch tensors.

  • BackendOps is the abstract operation surface.

  • NumpyOps is always available and uses NumPy plus SciPy sparse.

  • JaxOps is optional and follows JAX tracing and dtype configuration.

  • TorchOps is optional and follows PyTorch tensor, dtype, device, and autograd semantics.

  • CuPyOps is optional and follows CuPy/CuPy sparse semantics.

Choosing a backend#

Users normally choose a backend by constructing a Context:

import numpy as np
import spacecore as sc

ctx = sc.Context(sc.NumpyOps(), dtype=np.float64)

Backend portability means SpaceCore calls the same abstract operations and keeps the same space/operator model. It does not hide backend-specific dtype defaults, optional dependency availability, device placement, sparse support, tracing, or autograd behavior.

Autodoc#

class spacecore.backend.BackendOps[source]#

Bases: ABC

Public numerical contract for SpaceCore backends.

Common dense-array operations delegate to the backend’s Array API-compatible xp namespace. Subclasses provide backend-specific sparse conversion, dtype policy, indexing mutation, control flow, device/autograd behavior, and operations not covered by the Array API.

property family: str#

Backend family identifier.

property allow_sparse: bool#

Whether this backend supports sparse arrays.

property has_native_vmap: bool#

Whether vmap is implemented by the backend rather than a Python loop.

free_memory_bytes()[source]#

Return currently free device memory in bytes, or None if unknown.

The kernel dispatcher (ADR-016) uses this to gate materializing fast paths — those that allocate more than O(1) extra memory — against a memory budget before selecting them. The base implementation returns None (unknown): a backend that can cheaply query free memory (e.g. a GPU runtime) overrides this. When the budget is unknown the dispatcher treats any cost-carrying spec as unaffordable, so reporting None is always safe and never routes to a materializing kernel.

Return type:

int | None

abstract property dense_array: Type[Any]#

Dense array type accepted by this backend.

abstract property sparse_array: Tuple[Type[Any], ...] | None#

Sparse array types accepted by this backend, or None.

abstractmethod sanitize_dtype(dtype)[source]#

Normalize a dtype specifier for this backend.

Parameters:

dtype (Any | None)

Return type:

Any

is_dense(x)[source]#

Return whether x is a dense array for this backend.

Parameters:

x (Any)

Return type:

bool

is_sparse(x)[source]#

Return whether x is a sparse array for this backend.

Parameters:

x (Any)

Return type:

bool

is_array(x)[source]#

Return whether x is any array for this backend.

Parameters:

x (Any)

Return type:

bool

abstractmethod assparse(x, *, dtype=None)[source]#

Convert input to a backend sparse array.

Parameters:
  • x (Any)

  • dtype (Any | None)

Return type:

SparseArray

abstractmethod sparse_matmul(a, b)[source]#

Multiply a sparse array by a dense array.

Parameters:
  • a (SparseArray)

  • b (DenseArray)

Return type:

DenseArray

abstractmethod logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False)[source]#

Compute a stable log-sum-exp reduction.

Parameters:
  • a (DenseArray)

  • axis (int | Sequence[int] | None)

  • b (DenseArray | None)

  • keepdims (bool)

  • return_sign (bool)

Return type:

DenseArray | Tuple[DenseArray, DenseArray]

index_set(x, index, values, *, copy=True)[source]#

Return x with x[index] set to values.

With copy=True a mutable copy is updated and returned; with copy=False x is mutated in place. Immutable backends override this method.

Parameters:
  • x (DenseArray)

  • index (int | slice | Any | Tuple[Any, ...])

  • values (ArrayLike)

  • copy (bool)

Return type:

DenseArray

index_add(x, index, values, *, copy=True)[source]#

Return x with values accumulated into x[index].

With copy=True a mutable copy is updated and returned; with copy=False x is mutated in place. Immutable backends override this method.

Parameters:
  • x (DenseArray)

  • index (int | slice | Any | Tuple[Any, ...])

  • values (DenseArray)

  • copy (bool)

Return type:

DenseArray

abstractmethod ix_(*args)[source]#

Build open-mesh index arrays.

Parameters:

args (Any)

Return type:

Any

abstractmethod fori_loop(lower, upper, body_fun, init_val)[source]#

Run a counted loop primitive.

Parameters:
  • lower (int)

  • upper (int)

  • body_fun (Callable[[int, T], T])

  • init_val (T)

Return type:

T

abstractmethod while_loop(cond_fun, body_fun, init_val)[source]#

Run a while-loop primitive.

Parameters:
  • cond_fun (Callable[[T], bool])

  • body_fun (Callable[[T], T])

  • init_val (T)

Return type:

T

abstractmethod scan(f, init, xs, length=None, reverse=False, unroll=1)[source]#

Run a scan primitive.

Parameters:
  • f (Callable[[Carry, X], Tuple[Carry, Y]])

  • init (Carry)

  • xs (X)

  • length (int | None)

  • reverse (bool)

  • unroll (int)

Return type:

Tuple[Carry, Y]

abstractmethod cond(pred, true_fun, false_fun, *operands)[source]#

Run backend-compatible conditional branch selection.

Parameters:
  • pred (bool)

  • true_fun (Callable[[T], R])

  • false_fun (Callable[[T], R])

  • operands (Any)

Return type:

R

abstractmethod allclose_sparse(a, b, rtol=1e-05, atol=1e-08)[source]#

Compare sparse arrays elementwise within tolerances.

Parameters:
  • a (SparseArray)

  • b (SparseArray)

  • rtol (float)

  • atol (float)

Return type:

bool

property inf: DenseArray#

Positive infinity as a cached backend scalar.

property nan: DenseArray#

NaN as a cached backend scalar.

property pi: DenseArray#

Pi as a cached backend scalar.

property e: DenseArray#

Euler’s number as a cached backend scalar.

eps(dtype)[source]#

Machine epsilon for dtype.

Parameters:

dtype (Any)

Return type:

float

is_complex_dtype(dtype)[source]#

Return whether dtype is a complex floating type.

Parameters:

dtype (Any) – Backend or portable dtype specifier to inspect.

Returns:

True when dtype represents complex floating values.

Return type:

bool

real_dtype(dtype)[source]#

Return the real floating dtype with the same precision as dtype.

Parameters:

dtype (Any) – Backend or portable dtype specifier.

Returns:

dtype itself when it is already real-valued; otherwise float32 for complex64 and float64 for complex128.

Return type:

DType

get_dtype(x)[source]#

Return x.dtype after verifying x is a backend array.

Parameters:

x (Any)

Return type:

Any

shape(x)[source]#

Return x.shape as a tuple.

Parameters:

x (Any)

Return type:

tuple[int, …]

ndim(x)[source]#

Return the number of dimensions of x.

Parameters:

x (Any)

Return type:

int

size(x)[source]#

Return the total number of elements in x.

Parameters:

x (Any)

Return type:

int

asarray(x, dtype=None, **backend_kwargs)[source]#

Convert input to a dense backend array (delegates to xp.asarray).

Parameters:
  • x (Any)

  • dtype (Any | None)

  • backend_kwargs (Any)

Return type:

DenseArray

astype(x, dtype, **backend_kwargs)[source]#

Cast x to dtype, returning x unchanged when dtype is None.

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

  • backend_kwargs (Any)

Return type:

DenseArray

empty(shape, dtype=None)[source]#

Create an uninitialized array (delegates to xp.empty).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

zeros(shape, dtype=None)[source]#

Create a zero-filled array (delegates to xp.zeros).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

ones(shape, dtype=None)[source]#

Create a one-filled array (delegates to xp.ones).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

zeros_like(x, dtype=None)[source]#

Create a zero-filled array like x (delegates to xp.zeros_like).

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

Return type:

DenseArray

ones_like(x, dtype=None)[source]#

Create a one-filled array like x (delegates to xp.ones_like).

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

Return type:

DenseArray

full_like(x, value, dtype=None)[source]#

Create a value-filled array like x (delegates to xp.full_like).

Parameters:
  • x (DenseArray)

  • value (Any)

  • dtype (Any | None)

Return type:

DenseArray

arange(start, stop=None, step=None, dtype=None)[source]#

Create an evenly spaced range (delegates to xp.arange).

Parameters:
  • start (int)

  • stop (int | None)

  • step (int | None)

  • dtype (Any | None)

Return type:

DenseArray

full(shape, fill_value, dtype=None)[source]#

Create a value-filled array (delegates to xp.full).

Parameters:
  • shape (Tuple[int, ...])

  • fill_value (Any)

  • dtype (Any | None)

Return type:

DenseArray

eye(n, m=None, dtype=None)[source]#

Create an identity-like matrix (delegates to xp.eye).

Parameters:
  • n (int)

  • m (int | None)

  • dtype (Any | None)

Return type:

DenseArray

ravel(x)[source]#

Flatten x to one dimension.

Parameters:

x (DenseArray)

Return type:

DenseArray

reshape(x, shape)[source]#

Reshape x (delegates to xp.reshape).

Parameters:
  • x (DenseArray)

  • shape (Tuple[int, ...] | int)

Return type:

DenseArray

transpose(x, axes=None)[source]#

Permute dimensions of x.

Parameters:
  • x (DenseArray)

  • axes (Sequence[int] | None)

Return type:

DenseArray

swapaxes(x, axis1, axis2)[source]#

Swap two axes of x.

Parameters:
  • x (DenseArray)

  • axis1 (int)

  • axis2 (int)

Return type:

DenseArray

broadcast_to(x, shape)[source]#

Broadcast x to shape (delegates to xp.broadcast_to).

Parameters:
  • x (DenseArray)

  • shape (Tuple[int, ...])

Return type:

DenseArray

expand_dims(x, axis)[source]#

Insert singleton dimensions into x.

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int])

Return type:

DenseArray

squeeze(x, axis=None)[source]#

Remove singleton dimensions from x.

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

Return type:

DenseArray

moveaxis(x, source, destination)[source]#

Move axes of x to new positions.

Parameters:
  • x (DenseArray)

  • source (int | Sequence[int])

  • destination (int | Sequence[int])

Return type:

DenseArray

stack(arrays, axis=0)[source]#

Stack arrays along a new axis (delegates to xp.stack).

Parameters:
  • arrays (Sequence[DenseArray])

  • axis (int)

Return type:

DenseArray

hstack(arrays)[source]#

Stack arrays horizontally / column-wise (delegates to xp.hstack).

Concatenates along axis 0 for 1-D inputs and along axis 1 otherwise, matching NumPy hstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

vstack(arrays)[source]#

Stack arrays vertically / row-wise (delegates to xp.vstack).

Inputs are promoted to at least 2-D and concatenated along axis 0, matching NumPy vstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

dstack(arrays)[source]#

Stack arrays depth-wise along the third axis (delegates to xp.dstack).

Inputs are promoted to at least 3-D and concatenated along axis 2, matching NumPy dstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

column_stack(arrays)[source]#

Stack 1-D arrays as columns into a 2-D array (delegates to xp.column_stack).

1-D inputs become columns; higher-dimensional inputs are concatenated along axis 1, matching NumPy column_stack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

vmap(fn, in_axes=0, out_axes=0)[source]#

Vectorize fn over array axes using a Python-loop fallback.

Parameters:
  • fn (Callable)

  • in_axes (int | Sequence[int | None] | None)

  • out_axes (int | Sequence[int | None] | None)

Return type:

Callable

vectorize(pyfunc, *, excluded=None, signature=None)[source]#

Vectorize a scalar Python function over its array arguments.

Returns a callable that applies pyfunc elementwise, broadcasting the array arguments against one another and preserving the broadcast shape. Mirrors numpy.vectorize().

Parameters:
  • pyfunc (Callable) – Function called on scalar elements of the (broadcast) inputs.

  • excluded (Sequence[int] | None) – Positional argument indices passed through to pyfunc unchanged instead of being vectorized.

  • signature (str | None) – Generalized-ufunc signature (e.g. "(n),(n)->()"). Supported only on backends that provide a native vectorize.

Return type:

Callable

Notes

Delegates to the backend’s native vectorize (NumPy, JAX, CuPy) when available; otherwise applies a portable Python-loop fallback. The fallback does not support signature.

conj(x)[source]#

Complex conjugate of x (delegates to xp.conj).

Parameters:

x (DenseArray)

Return type:

DenseArray

real(x)[source]#

Real component of x (delegates to xp.real).

Parameters:

x (DenseArray)

Return type:

DenseArray

imag(x)[source]#

Imaginary component of x (delegates to xp.imag).

Parameters:

x (DenseArray)

Return type:

DenseArray

abs(x)[source]#

Absolute value of x (delegates to xp.abs).

Parameters:

x (DenseArray)

Return type:

DenseArray

sign(x)[source]#

Elementwise sign of x (delegates to xp.sign).

Parameters:

x (DenseArray)

Return type:

DenseArray

sqrt(x)[source]#

Elementwise square root of x (delegates to xp.sqrt).

Parameters:

x (DenseArray)

Return type:

DenseArray

sum(x, axis=None, keepdims=False, dtype=None)[source]#

Sum over given axes (delegates to xp.sum).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

Return type:

DenseArray

mean(x, axis=None, keepdims=False)[source]#

Mean over given axes (delegates to xp.mean).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

min(x, axis=None, keepdims=False)[source]#

Minimum over given axes (delegates to xp.min).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

max(x, axis=None, keepdims=False)[source]#

Maximum over given axes (delegates to xp.max).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

prod(x, axis=None, keepdims=False, dtype=None)[source]#

Product over given axes (delegates to xp.prod).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

Return type:

DenseArray

trace(x)[source]#

Trace of a matrix (delegates to xp.trace when available).

Parameters:

x (DenseArray)

Return type:

DenseArray

argsort(x, axis=-1)[source]#

Return indices that sort x along an axis.

Parameters:
  • x (DenseArray)

  • axis (int)

Return type:

DenseArray

sort(x, axis=-1)[source]#

Sort x along an axis (delegates to xp.sort).

Parameters:
  • x (DenseArray)

  • axis (int)

Return type:

DenseArray

argmin(x, axis=None, keepdims=False)[source]#

Return indices of minima along an axis.

Parameters:
  • x (DenseArray)

  • axis (int | None)

  • keepdims (bool)

Return type:

DenseArray

argmax(x, axis=None, keepdims=False)[source]#

Return indices of maxima along an axis.

Parameters:
  • x (DenseArray)

  • axis (int | None)

  • keepdims (bool)

Return type:

DenseArray

vdot(x, y)[source]#

Return sum(conj(x) * y) over flattened inputs.

Matches NumPy, JAX, and Torch vdot semantics. DenseLinOp.rapply relies on this convention for complex inputs.

Parameters:
  • x (DenseArray)

  • y (DenseArray)

Return type:

DenseArray

matmul(a, b, backend_kwargs=None)[source]#

Matrix product (delegates to xp.matmul).

Parameters:
  • a (DenseArray)

  • b (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

kron(a, b)[source]#

Kronecker product (delegates to xp.kron).

Parameters:
  • a (DenseArray)

  • b (DenseArray)

Return type:

DenseArray

einsum(subscripts, *operands)[source]#

Einstein summation (delegates to xp.einsum).

Parameters:
  • subscripts (str)

  • operands (DenseArray)

Return type:

DenseArray

eigh(x, backend_kwargs=None)[source]#

Eigenpairs of a Hermitian dense matrix (delegates to xp.linalg.eigh).

Parameters:
  • x (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseArray, DenseArray]

norm(x, ord=None, axis=None, keepdims=False)[source]#

Vector or matrix norm (delegates to xp.linalg.norm).

Parameters:
  • x (DenseArray)

  • ord (int | str | None)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

solve(A, b, backend_kwargs=None)[source]#

Solve a dense linear system (delegates to xp.linalg.solve).

Parameters:
  • A (DenseArray)

  • b (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

eigvalsh(A, backend_kwargs=None)[source]#

Eigenvalues of a Hermitian dense matrix (delegates to xp.linalg.eigvalsh).

Parameters:
  • A (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

svd(A, full_matrices=True, backend_kwargs=None)[source]#

Singular value decomposition (delegates to xp.linalg.svd).

Parameters:
  • A (DenseArray)

  • full_matrices (bool)

  • backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseArray, DenseArray, DenseArray]

cholesky(A, backend_kwargs=None)[source]#

Cholesky factorization (delegates to xp.linalg.cholesky).

Parameters:
  • A (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

exp(x)[source]#

Elementwise exponential (delegates to xp.exp).

Parameters:

x (DenseArray)

Return type:

DenseArray

log(x)[source]#

Elementwise natural logarithm (delegates to xp.log).

Parameters:

x (DenseArray)

Return type:

DenseArray

where(condition, x, y)[source]#

Select between x and y by condition (delegates to xp.where).

Parameters:
  • condition (DenseArray | bool)

  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

maximum(x, y)[source]#

Elementwise maximum (delegates to xp.maximum).

Parameters:
  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

minimum(x, y)[source]#

Elementwise minimum (delegates to xp.minimum).

Parameters:
  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

clip(x, a_min, a_max)[source]#

Clip x into [a_min, a_max] (delegates to xp.clip).

Parameters:
  • x (DenseArray)

  • a_min (ArrayLike)

  • a_max (ArrayLike)

Return type:

DenseArray

isfinite(x)[source]#

Elementwise finite check (delegates to xp.isfinite).

Parameters:

x (DenseArray)

Return type:

DenseArray

isnan(x)[source]#

Elementwise NaN check (delegates to xp.isnan).

Parameters:

x (DenseArray)

Return type:

DenseArray

concatenate(arrays, axis=0, dtype=None)[source]#

Concatenate arrays along an existing axis (delegates to xp.concat).

Parameters:
  • arrays (Sequence[DenseArray])

  • axis (int)

  • dtype (Any | None)

Return type:

DenseArray

take(x, indices, axis=None)[source]#

Take entries from x by integer indices (delegates to xp.take).

Parameters:
  • x (DenseArray)

  • indices (DenseArray)

  • axis (int | None)

Return type:

DenseArray

diag(x)[source]#

Extract or construct a diagonal (delegates to xp.diag).

Parameters:

x (DenseArray)

Return type:

DenseArray

diagonal(x)[source]#

Return the main diagonal of x (delegates to xp.diagonal).

Parameters:

x (DenseArray)

Return type:

DenseArray

tril(x)[source]#

Lower triangle of x (delegates to xp.tril).

Parameters:

x (DenseArray)

Return type:

DenseArray

triu(x)[source]#

Upper triangle of x (delegates to xp.triu).

Parameters:

x (DenseArray)

Return type:

DenseArray

allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)[source]#

Return whether dense arrays are close within tolerances.

Parameters:
  • a (DenseArray)

  • b (DenseArray)

  • rtol (float)

  • atol (float)

  • equal_nan (bool)

Return type:

bool

class spacecore.backend.NumpyOps[source]#

Bases: EagerControlFlowMixin, BackendOps

BackendOps implementation for the NumPy/SciPy ecosystem.

This backend uses NumPy for dense array operations and SciPy for sparse array operations.

Dense arrays

numpy.ndarray

Sparse arrays

scipy.sparse.spmatrix scipy.sparse.sparray

Most methods mirror the corresponding NumPy or SciPy signatures and
delegate directly to NumPy/SciPy implementations. Backend-specific
behavior, dtype promotion, broadcasting, memory layout, and error modes
therefore follow NumPy/SciPy semantics.
Backend handles
  • np : module NumPy module stored on the class and available through instances as ops.np. Advanced users may use it when SpaceCore’s portable API does not expose a required NumPy feature.

  • sp : module SciPy module stored on the class and available through instances as ops.sp. Advanced users may use it for SciPy-specific functionality.

Notes

Code intended to remain backend-portable should prefer BackendOps methods. Direct use of ops.np or ops.sp is an explicit NumPy/SciPy-specific escape hatch.

NumPy’s device keyword is present for Array API interoperability. When supplied, it must be “cpu” or None; see the corresponding NumPy documentation for each method.

np: Any = <module 'numpy' from '/opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/numpy/__init__.py'>#
sp: Any = <module 'scipy' from '/opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/scipy/__init__.py'>#
free_memory_bytes()[source]#

Return available system RAM in bytes.

For a CPU backend the “device” is host memory, so this reports the available physical RAM via psutil, a required dependency. The kernel dispatcher (ADR-016) uses it as the memory budget that gates a materializing fast path.

Return type:

int | None

property dense_array: Type[Any]#

Dense array type using NumPy.

Return type:

Concrete dense array class accepted by this backend.

See:

https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html

property sparse_array: Tuple[Type[Any], ...]#

Sparse array type tuple using SciPy.

Return type:

Concrete sparse array classes accepted by this backend, or None.

See:

https://docs.scipy.org/doc/scipy/reference/sparse.html

sanitize_dtype(dtype)[source]#

Normalize a dtype specifier using NumPy.

Input:

dtype: Optional dtype requested by SpaceCore or the caller.

Output:

Backend dtype object accepted by array constructors.

See:

https://numpy.org/doc/stable/reference/generated/numpy.dtype.html

Parameters:

dtype (Any | None)

Return type:

Any

assparse(x, *, format='csr', dtype=None)[source]#

Convert input to a sparse array using SciPy.

Input:

x: Dense, sparse, or array-like input plus sparse-format options.

Output:

Sparse backend array.

See:

https://docs.scipy.org/doc/scipy/reference/sparse.html

Backend-specific notes:

SpaceCore currently converts dense inputs to 2-D SciPy sparse matrices in the requested format.

Parameters:
  • x (Any)

  • format (Literal['csr', 'csc', 'coo'])

  • dtype (Any | None)

Return type:

SparseArray

sparse_matmul(a, b)[source]#

Multiply sparse and dense arrays using SciPy.

Input:

a: Sparse backend array; b: Dense backend array.

Output:

Dense backend array containing the product.

See:

https://docs.scipy.org/doc/scipy/reference/sparse.html

Backend-specific notes:

Uses SciPy sparse multiplication before returning a dense NumPy result when applicable.

Parameters:
  • a (SparseArray)

  • b (DenseArray)

Return type:

DenseArray

logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False)[source]#

Compute a stable log-sum-exp reduction using SciPy.

Input:

a: Dense backend array; axis, weights, and sign options control the reduction.

Output:

Dense backend array or tuple containing log-sum-exp results.

See:

https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.logsumexp.html

Parameters:
  • a (DenseArray)

  • axis (int | Sequence[int] | None)

  • b (DenseArray | None)

  • keepdims (bool)

  • return_sign (bool)

Return type:

DenseArray | Tuple[DenseArray, DenseArray]

ix_(*args)[source]#

Build open mesh index arrays using NumPy.

Input:

args: One-dimensional index arrays or sequences.

Output:

Tuple of dense backend arrays usable for open-mesh indexing.

See:

https://numpy.org/doc/stable/reference/generated/numpy.ix\_.html

Parameters:

args (Any)

Return type:

Any

allclose_sparse(a, b, rtol=1e-05, atol=1e-08)[source]#

Compare sparse arrays elementwise within tolerances using SciPy.

Input:

a, b: Sparse backend arrays; rtol and atol configure comparison.

Output:

Boolean indicating whether sparse arrays are close.

See:

https://docs.scipy.org/doc/scipy/reference/sparse.html

Backend-specific notes:

Sparse inputs are converted to CSR and compared by logical sparse difference.

Parameters:
  • a (SparseArray)

  • b (SparseArray)

  • rtol (float)

  • atol (float)

Return type:

bool

abs(x)#

Absolute value of x (delegates to xp.abs).

Parameters:

x (DenseArray)

Return type:

DenseArray

allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)#

Return whether dense arrays are close within tolerances.

Parameters:
  • a (DenseArray)

  • b (DenseArray)

  • rtol (float)

  • atol (float)

  • equal_nan (bool)

Return type:

bool

property allow_sparse: bool#

Whether this backend supports sparse arrays.

arange(start, stop=None, step=None, dtype=None)#

Create an evenly spaced range (delegates to xp.arange).

Parameters:
  • start (int)

  • stop (int | None)

  • step (int | None)

  • dtype (Any | None)

Return type:

DenseArray

argmax(x, axis=None, keepdims=False)#

Return indices of maxima along an axis.

Parameters:
  • x (DenseArray)

  • axis (int | None)

  • keepdims (bool)

Return type:

DenseArray

argmin(x, axis=None, keepdims=False)#

Return indices of minima along an axis.

Parameters:
  • x (DenseArray)

  • axis (int | None)

  • keepdims (bool)

Return type:

DenseArray

argsort(x, axis=-1)#

Return indices that sort x along an axis.

Parameters:
  • x (DenseArray)

  • axis (int)

Return type:

DenseArray

asarray(x, dtype=None, **backend_kwargs)#

Convert input to a dense backend array (delegates to xp.asarray).

Parameters:
  • x (Any)

  • dtype (Any | None)

  • backend_kwargs (Any)

Return type:

DenseArray

astype(x, dtype, **backend_kwargs)#

Cast x to dtype, returning x unchanged when dtype is None.

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

  • backend_kwargs (Any)

Return type:

DenseArray

broadcast_to(x, shape)#

Broadcast x to shape (delegates to xp.broadcast_to).

Parameters:
  • x (DenseArray)

  • shape (Tuple[int, ...])

Return type:

DenseArray

cholesky(A, backend_kwargs=None)#

Cholesky factorization (delegates to xp.linalg.cholesky).

Parameters:
  • A (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

clip(x, a_min, a_max)#

Clip x into [a_min, a_max] (delegates to xp.clip).

Parameters:
  • x (DenseArray)

  • a_min (ArrayLike)

  • a_max (ArrayLike)

Return type:

DenseArray

column_stack(arrays)#

Stack 1-D arrays as columns into a 2-D array (delegates to xp.column_stack).

1-D inputs become columns; higher-dimensional inputs are concatenated along axis 1, matching NumPy column_stack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

concatenate(arrays, axis=0, dtype=None)#

Concatenate arrays along an existing axis (delegates to xp.concat).

Parameters:
  • arrays (Sequence[DenseArray])

  • axis (int)

  • dtype (Any | None)

Return type:

DenseArray

cond(pred, true_fun, false_fun, *operands)#

Select a branch eagerly using Python truth-value conversion.

Parameters:
  • pred (bool)

  • true_fun (Callable[[T], R])

  • false_fun (Callable[[T], R])

  • operands (Any)

Return type:

R

conj(x)#

Complex conjugate of x (delegates to xp.conj).

Parameters:

x (DenseArray)

Return type:

DenseArray

diag(x)#

Extract or construct a diagonal (delegates to xp.diag).

Parameters:

x (DenseArray)

Return type:

DenseArray

diagonal(x)#

Return the main diagonal of x (delegates to xp.diagonal).

Parameters:

x (DenseArray)

Return type:

DenseArray

dstack(arrays)#

Stack arrays depth-wise along the third axis (delegates to xp.dstack).

Inputs are promoted to at least 3-D and concatenated along axis 2, matching NumPy dstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

property e: DenseArray#

Euler’s number as a cached backend scalar.

eigh(x, backend_kwargs=None)#

Eigenpairs of a Hermitian dense matrix (delegates to xp.linalg.eigh).

Parameters:
  • x (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseArray, DenseArray]

eigvalsh(A, backend_kwargs=None)#

Eigenvalues of a Hermitian dense matrix (delegates to xp.linalg.eigvalsh).

Parameters:
  • A (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

einsum(subscripts, *operands)#

Einstein summation (delegates to xp.einsum).

Parameters:
  • subscripts (str)

  • operands (DenseArray)

Return type:

DenseArray

empty(shape, dtype=None)#

Create an uninitialized array (delegates to xp.empty).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

eps(dtype)#

Machine epsilon for dtype.

Parameters:

dtype (Any)

Return type:

float

exp(x)#

Elementwise exponential (delegates to xp.exp).

Parameters:

x (DenseArray)

Return type:

DenseArray

expand_dims(x, axis)#

Insert singleton dimensions into x.

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int])

Return type:

DenseArray

eye(n, m=None, dtype=None)#

Create an identity-like matrix (delegates to xp.eye).

Parameters:
  • n (int)

  • m (int | None)

  • dtype (Any | None)

Return type:

DenseArray

property family: str#

Backend family identifier.

fori_loop(lower, upper, body_fun, init_val)#

Run a counted loop eagerly as a Python for loop.

This executes without JAX tracing semantics; unroll-style parameters are not accepted because eager backends gain nothing from them.

Parameters:
  • lower (int)

  • upper (int)

  • body_fun (Callable[[int, T], T])

  • init_val (T)

Return type:

T

full(shape, fill_value, dtype=None)#

Create a value-filled array (delegates to xp.full).

Parameters:
  • shape (Tuple[int, ...])

  • fill_value (Any)

  • dtype (Any | None)

Return type:

DenseArray

full_like(x, value, dtype=None)#

Create a value-filled array like x (delegates to xp.full_like).

Parameters:
  • x (DenseArray)

  • value (Any)

  • dtype (Any | None)

Return type:

DenseArray

get_dtype(x)#

Return x.dtype after verifying x is a backend array.

Parameters:

x (Any)

Return type:

Any

property has_native_vmap: bool#

Whether vmap is implemented by the backend rather than a Python loop.

hstack(arrays)#

Stack arrays horizontally / column-wise (delegates to xp.hstack).

Concatenates along axis 0 for 1-D inputs and along axis 1 otherwise, matching NumPy hstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

imag(x)#

Imaginary component of x (delegates to xp.imag).

Parameters:

x (DenseArray)

Return type:

DenseArray

index_add(x, index, values, *, copy=True)#

Return x with values accumulated into x[index].

With copy=True a mutable copy is updated and returned; with copy=False x is mutated in place. Immutable backends override this method.

Parameters:
  • x (DenseArray)

  • index (int | slice | Any | Tuple[Any, ...])

  • values (DenseArray)

  • copy (bool)

Return type:

DenseArray

index_set(x, index, values, *, copy=True)#

Return x with x[index] set to values.

With copy=True a mutable copy is updated and returned; with copy=False x is mutated in place. Immutable backends override this method.

Parameters:
  • x (DenseArray)

  • index (int | slice | Any | Tuple[Any, ...])

  • values (ArrayLike)

  • copy (bool)

Return type:

DenseArray

property inf: DenseArray#

Positive infinity as a cached backend scalar.

is_array(x)#

Return whether x is any array for this backend.

Parameters:

x (Any)

Return type:

bool

is_complex_dtype(dtype)#

Return whether dtype is a complex floating type.

Parameters:

dtype (Any) – Backend or portable dtype specifier to inspect.

Returns:

True when dtype represents complex floating values.

Return type:

bool

is_dense(x)#

Return whether x is a dense array for this backend.

Parameters:

x (Any)

Return type:

bool

is_sparse(x)#

Return whether x is a sparse array for this backend.

Parameters:

x (Any)

Return type:

bool

isfinite(x)#

Elementwise finite check (delegates to xp.isfinite).

Parameters:

x (DenseArray)

Return type:

DenseArray

isnan(x)#

Elementwise NaN check (delegates to xp.isnan).

Parameters:

x (DenseArray)

Return type:

DenseArray

kron(a, b)#

Kronecker product (delegates to xp.kron).

Parameters:
  • a (DenseArray)

  • b (DenseArray)

Return type:

DenseArray

log(x)#

Elementwise natural logarithm (delegates to xp.log).

Parameters:

x (DenseArray)

Return type:

DenseArray

matmul(a, b, backend_kwargs=None)#

Matrix product (delegates to xp.matmul).

Parameters:
  • a (DenseArray)

  • b (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

max(x, axis=None, keepdims=False)#

Maximum over given axes (delegates to xp.max).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

maximum(x, y)#

Elementwise maximum (delegates to xp.maximum).

Parameters:
  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

mean(x, axis=None, keepdims=False)#

Mean over given axes (delegates to xp.mean).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

min(x, axis=None, keepdims=False)#

Minimum over given axes (delegates to xp.min).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

minimum(x, y)#

Elementwise minimum (delegates to xp.minimum).

Parameters:
  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

moveaxis(x, source, destination)#

Move axes of x to new positions.

Parameters:
  • x (DenseArray)

  • source (int | Sequence[int])

  • destination (int | Sequence[int])

Return type:

DenseArray

property nan: DenseArray#

NaN as a cached backend scalar.

ndim(x)#

Return the number of dimensions of x.

Parameters:

x (Any)

Return type:

int

norm(x, ord=None, axis=None, keepdims=False)#

Vector or matrix norm (delegates to xp.linalg.norm).

Parameters:
  • x (DenseArray)

  • ord (int | str | None)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

ones(shape, dtype=None)#

Create a one-filled array (delegates to xp.ones).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

ones_like(x, dtype=None)#

Create a one-filled array like x (delegates to xp.ones_like).

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

Return type:

DenseArray

property pi: DenseArray#

Pi as a cached backend scalar.

prod(x, axis=None, keepdims=False, dtype=None)#

Product over given axes (delegates to xp.prod).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

Return type:

DenseArray

ravel(x)#

Flatten x to one dimension.

Parameters:

x (DenseArray)

Return type:

DenseArray

real(x)#

Real component of x (delegates to xp.real).

Parameters:

x (DenseArray)

Return type:

DenseArray

real_dtype(dtype)#

Return the real floating dtype with the same precision as dtype.

Parameters:

dtype (Any) – Backend or portable dtype specifier.

Returns:

dtype itself when it is already real-valued; otherwise float32 for complex64 and float64 for complex128.

Return type:

DType

reshape(x, shape)#

Reshape x (delegates to xp.reshape).

Parameters:
  • x (DenseArray)

  • shape (Tuple[int, ...] | int)

Return type:

DenseArray

scan(f, init, xs, length=None, reverse=False, unroll=1)#

Run a scan primitive eagerly, stacking per-step outputs at the end.

unroll is accepted only for API parity with jax.lax.scan. When xs is None an explicit length is required; otherwise the length is inferred from the leading axis of the first leaf of xs.

Parameters:
  • f (Callable[[Carry, X], Tuple[Carry, Y]])

  • init (Carry)

  • xs (X)

  • length (int | None)

  • reverse (bool)

  • unroll (int)

Return type:

Tuple[Carry, Y]

shape(x)#

Return x.shape as a tuple.

Parameters:

x (Any)

Return type:

tuple[int, …]

sign(x)#

Elementwise sign of x (delegates to xp.sign).

Parameters:

x (DenseArray)

Return type:

DenseArray

size(x)#

Return the total number of elements in x.

Parameters:

x (Any)

Return type:

int

solve(A, b, backend_kwargs=None)#

Solve a dense linear system (delegates to xp.linalg.solve).

Parameters:
  • A (DenseArray)

  • b (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

sort(x, axis=-1)#

Sort x along an axis (delegates to xp.sort).

Parameters:
  • x (DenseArray)

  • axis (int)

Return type:

DenseArray

sqrt(x)#

Elementwise square root of x (delegates to xp.sqrt).

Parameters:

x (DenseArray)

Return type:

DenseArray

squeeze(x, axis=None)#

Remove singleton dimensions from x.

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

Return type:

DenseArray

stack(arrays, axis=0)#

Stack arrays along a new axis (delegates to xp.stack).

Parameters:
  • arrays (Sequence[DenseArray])

  • axis (int)

Return type:

DenseArray

sum(x, axis=None, keepdims=False, dtype=None)#

Sum over given axes (delegates to xp.sum).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

Return type:

DenseArray

svd(A, full_matrices=True, backend_kwargs=None)#

Singular value decomposition (delegates to xp.linalg.svd).

Parameters:
  • A (DenseArray)

  • full_matrices (bool)

  • backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseArray, DenseArray, DenseArray]

swapaxes(x, axis1, axis2)#

Swap two axes of x.

Parameters:
  • x (DenseArray)

  • axis1 (int)

  • axis2 (int)

Return type:

DenseArray

take(x, indices, axis=None)#

Take entries from x by integer indices (delegates to xp.take).

Parameters:
  • x (DenseArray)

  • indices (DenseArray)

  • axis (int | None)

Return type:

DenseArray

trace(x)#

Trace of a matrix (delegates to xp.trace when available).

Parameters:

x (DenseArray)

Return type:

DenseArray

transpose(x, axes=None)#

Permute dimensions of x.

Parameters:
  • x (DenseArray)

  • axes (Sequence[int] | None)

Return type:

DenseArray

tril(x)#

Lower triangle of x (delegates to xp.tril).

Parameters:

x (DenseArray)

Return type:

DenseArray

triu(x)#

Upper triangle of x (delegates to xp.triu).

Parameters:

x (DenseArray)

Return type:

DenseArray

vdot(x, y)#

Return sum(conj(x) * y) over flattened inputs.

Matches NumPy, JAX, and Torch vdot semantics. DenseLinOp.rapply relies on this convention for complex inputs.

Parameters:
  • x (DenseArray)

  • y (DenseArray)

Return type:

DenseArray

vectorize(pyfunc, *, excluded=None, signature=None)#

Vectorize a scalar Python function over its array arguments.

Returns a callable that applies pyfunc elementwise, broadcasting the array arguments against one another and preserving the broadcast shape. Mirrors numpy.vectorize().

Parameters:
  • pyfunc (Callable) – Function called on scalar elements of the (broadcast) inputs.

  • excluded (Sequence[int] | None) – Positional argument indices passed through to pyfunc unchanged instead of being vectorized.

  • signature (str | None) – Generalized-ufunc signature (e.g. "(n),(n)->()"). Supported only on backends that provide a native vectorize.

Return type:

Callable

Notes

Delegates to the backend’s native vectorize (NumPy, JAX, CuPy) when available; otherwise applies a portable Python-loop fallback. The fallback does not support signature.

vmap(fn, in_axes=0, out_axes=0)#

Vectorize fn over array axes using a Python-loop fallback.

Parameters:
  • fn (Callable)

  • in_axes (int | Sequence[int | None] | None)

  • out_axes (int | Sequence[int | None] | None)

Return type:

Callable

vstack(arrays)#

Stack arrays vertically / row-wise (delegates to xp.vstack).

Inputs are promoted to at least 2-D and concatenated along axis 0, matching NumPy vstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

where(condition, x, y)#

Select between x and y by condition (delegates to xp.where).

Parameters:
  • condition (DenseArray | bool)

  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

while_loop(cond_fun, body_fun, init_val)#

Run a while-loop eagerly, converting the predicate to bool each step.

Parameters:
  • cond_fun (Callable[[T], bool])

  • body_fun (Callable[[T], T])

  • init_val (T)

Return type:

T

zeros(shape, dtype=None)#

Create a zero-filled array (delegates to xp.zeros).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

zeros_like(x, dtype=None)#

Create a zero-filled array like x (delegates to xp.zeros_like).

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

Return type:

DenseArray

class spacecore.backend.JaxOps[source]#

Bases: BackendOps

BackendOps implementation for the JAX ecosystem.

This backend uses JAX for dense array operations and JAX experimental sparse arrays for sparse operations.

Dense arrays

jax.Array

Sparse arrays

jax.experimental.sparse.BCOO jax.experimental.sparse.BCSR

Most methods mirror the corresponding JAX public API signatures and
delegate to `jax.numpy`, `jax.numpy.linalg`, `jax.scipy`, or
`jax.experimental.sparse`. Backend-specific behavior, tracing rules,
dtype canonicalization, device placement, sharding, and error modes
therefore follow JAX semantics.
Backend handles
  • jaxmodule

    JAX module stored on the class and available through instances as ops.jax. Advanced users may use it when SpaceCore’s portable API does not expose a required JAX feature.

  • jnpmodule

    jax.numpy module stored on the class and available through instances as ops.jnp.

  • jsparsemodule

    jax.experimental.sparse module stored on the class and available through instances as ops.jsparse.

Notes

Code intended to remain backend-portable should prefer BackendOps methods. Direct use of ops.jax, ops.jnp, or ops.jsparse is an explicit JAX-specific escape hatch.

Some parameters are accepted for JAX signature compatibility even when JAX ignores them. Array-creation routines may expose device and out_sharding for explicit placement or sharding.

sanitize_dtype(dtype)[source]#

Normalize a dtype specifier using JAX.

Input:

dtype: Optional dtype requested by SpaceCore or the caller.

Output:

Backend dtype object accepted by array constructors.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.numpy.dtype.html

Backend-specific notes:

SpaceCore rejects dtypes that JAX would silently canonicalize under the active x64 setting.

Parameters:

dtype (Any | None)

Return type:

Any

property dense_array: Type[Any]#

Dense array type using JAX.

Return type:

Concrete dense array class accepted by this backend.

See:

https://docs.jax.dev/en/latest/jax.Array.html

property sparse_array: Tuple[Type[Any], ...]#

Sparse array type tuple using JAX.

Return type:

Concrete sparse array classes accepted by this backend, or None.

See:

https://docs.jax.dev/en/latest/jax.experimental.sparse.html

assparse(x, *, format='bcoo', index_dtype=None, nse=None, dtype=None)[source]#

Convert input to a sparse array using JAX.

Input:

x: Dense, sparse, or array-like input plus sparse-format options.

Output:

Sparse backend array.

See:

https://docs.jax.dev/en/latest/jax.experimental.sparse.html

Backend-specific notes:

Dense inputs are converted with JAX sparse BCOO/BCSR constructors; SciPy sparse inputs use from_scipy_sparse.

Parameters:
  • x (Any)

  • format (Literal['bcoo', 'bcsr'])

  • index_dtype (Any | None)

  • nse (int | None)

  • dtype (Any | None)

Return type:

SparseArray

sparse_matmul(a, b)[source]#

Multiply sparse and dense arrays using JAX.

Input:

a: Sparse backend array; b: Dense backend array.

Output:

Dense backend array containing the product.

See:

https://docs.jax.dev/en/latest/jax.experimental.sparse.html

Backend-specific notes:

Uses JAX sparse matmul and returns a JAX array; sparse support remains experimental in JAX.

Parameters:
  • a (SparseArray)

  • b (DenseArray)

Return type:

DenseArray

vmap(fn, in_axes=0, out_axes=0)[source]#

Vectorize a function using jax.vmap.

Parameters:
  • fn (Callable)

  • in_axes (int | Sequence[int | None] | None)

  • out_axes (int | Sequence[int | None] | None)

Return type:

Callable

property has_native_vmap: bool#

Return True because JAX provides native vmap.

logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False, where=None)[source]#

Compute a stable log-sum-exp reduction using JAX.

Input:

a: Dense backend array; axis, weights, and sign options control the reduction.

Output:

Dense backend array or tuple containing log-sum-exp results.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.scipy.special.logsumexp.html

Parameters:
  • a (DenseArray)

  • axis (int | Sequence[int] | None)

  • b (DenseArray | None)

  • keepdims (bool)

  • return_sign (bool)

  • where (DenseArray | None)

Return type:

DenseArray | Tuple[DenseArray, DenseArray]

index_set(x, index, values, *, copy=True)[source]#

Set indexed values using JAX.

Input:

x: Dense backend array; index: Selection; values: Replacement values; copy controls mutation policy.

Output:

Dense backend array with indexed values set.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.Array.at.html

Backend-specific notes:

JAX arrays are immutable; copy=False raises NotImplementedError.

Parameters:
  • x (DenseArray)

  • index (int | slice | Any | Tuple[Any, ...])

  • values (ArrayLike)

  • copy (bool)

ix_(*args)[source]#

Build open mesh index arrays using JAX.

Input:

args: One-dimensional index arrays or sequences.

Output:

Tuple of dense backend arrays usable for open-mesh indexing.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ix\_.html

Parameters:

args (Any)

Return type:

Any

fori_loop(lower, upper, body_fun, init_val, *, unroll=None)[source]#

Run a counted loop primitive using JAX.

Input:

lower, upper: Loop bounds; body_fun: Loop body; init_val: Initial carry value.

Output:

Final carry value after loop execution.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.lax.fori_loop.html

Backend-specific notes:

Loop bounds and unroll behavior follow JAX tracing and compilation rules.

Parameters:
  • lower (int)

  • upper (int)

  • body_fun (Callable[[int, T], T])

  • init_val (T)

  • unroll (int | bool | None)

Return type:

T

while_loop(cond_fun, body_fun, init_val)[source]#

Run a while-loop primitive using JAX.

Input:

cond_fun: Loop condition; body_fun: Loop body; init_val: Initial carry value.

Output:

Final carry value after loop execution.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.lax.while_loop.html

Backend-specific notes:

Condition and body are staged according to JAX lax control-flow semantics.

Parameters:
  • cond_fun (Callable[[T], bool])

  • body_fun (Callable[[T], T])

  • init_val (T)

Return type:

T

scan(f, init, xs, length=None, reverse=False, unroll=1, _split_transpose=False)[source]#

Run a scan primitive using JAX.

Input:

f: Scan body; init: Initial carry; xs: Per-step inputs plus scan options.

Output:

Tuple of final carry and stacked outputs.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.lax.scan.html

Backend-specific notes:

Inputs and outputs may be pytrees and are staged according to JAX lax.scan semantics.

Parameters:
  • f (Callable[[Carry, X], Tuple[Carry, Y]])

  • init (Carry)

  • xs (X)

  • length (int | None)

  • reverse (bool)

  • unroll (int)

  • _split_transpose (bool)

Return type:

Tuple[Carry, Y]

cond(pred, true_fun, false_fun, *operands)[source]#

Run conditional branch selection using JAX.

Input:

pred: Predicate; true_fun and false_fun: Branch functions; operands: Branch inputs.

Output:

Result returned by the selected branch.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.lax.cond.html

Backend-specific notes:

Branches are staged according to JAX lax.cond semantics rather than Python eager branching.

Parameters:
  • pred (bool)

  • true_fun (Callable[[T], R])

  • false_fun (Callable[[T], R])

  • operands (Any)

Return type:

R

index_add(x, index, values, *, copy=True)[source]#

Add into indexed values using JAX.

Input:

x: Dense backend array; index: Selection; values: Values to add; copy controls mutation policy.

Output:

Dense backend array with indexed values incremented.

See:

https://docs.jax.dev/en/latest/_autosummary/jax.Array.at.html

Backend-specific notes:

JAX arrays are immutable; copy=False raises NotImplementedError and repeated indices follow JAX scatter-add semantics.

Parameters:
  • x (DenseArray)

  • index (int | slice | Any | Tuple[Any, ...])

  • values (DenseArray)

  • copy (bool)

allclose_sparse(a, b, rtol=1e-05, atol=1e-08)[source]#

Compare sparse arrays elementwise within tolerances using JAX.

Input:

a, b: Sparse backend arrays; rtol and atol configure comparison.

Output:

Boolean indicating whether sparse arrays are close.

See:

https://docs.jax.dev/en/latest/jax.experimental.sparse.html

Backend-specific notes:

SpaceCore converts JAX sparse arrays through SciPy sparse arrays for comparison.

Parameters:
  • a (SparseArray)

  • b (SparseArray)

  • rtol (float)

  • atol (float)

Return type:

bool

abs(x)#

Absolute value of x (delegates to xp.abs).

Parameters:

x (DenseArray)

Return type:

DenseArray

allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)#

Return whether dense arrays are close within tolerances.

Parameters:
  • a (DenseArray)

  • b (DenseArray)

  • rtol (float)

  • atol (float)

  • equal_nan (bool)

Return type:

bool

property allow_sparse: bool#

Whether this backend supports sparse arrays.

arange(start, stop=None, step=None, dtype=None)#

Create an evenly spaced range (delegates to xp.arange).

Parameters:
  • start (int)

  • stop (int | None)

  • step (int | None)

  • dtype (Any | None)

Return type:

DenseArray

argmax(x, axis=None, keepdims=False)#

Return indices of maxima along an axis.

Parameters:
  • x (DenseArray)

  • axis (int | None)

  • keepdims (bool)

Return type:

DenseArray

argmin(x, axis=None, keepdims=False)#

Return indices of minima along an axis.

Parameters:
  • x (DenseArray)

  • axis (int | None)

  • keepdims (bool)

Return type:

DenseArray

argsort(x, axis=-1)#

Return indices that sort x along an axis.

Parameters:
  • x (DenseArray)

  • axis (int)

Return type:

DenseArray

asarray(x, dtype=None, **backend_kwargs)#

Convert input to a dense backend array (delegates to xp.asarray).

Parameters:
  • x (Any)

  • dtype (Any | None)

  • backend_kwargs (Any)

Return type:

DenseArray

astype(x, dtype, **backend_kwargs)#

Cast x to dtype, returning x unchanged when dtype is None.

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

  • backend_kwargs (Any)

Return type:

DenseArray

broadcast_to(x, shape)#

Broadcast x to shape (delegates to xp.broadcast_to).

Parameters:
  • x (DenseArray)

  • shape (Tuple[int, ...])

Return type:

DenseArray

cholesky(A, backend_kwargs=None)#

Cholesky factorization (delegates to xp.linalg.cholesky).

Parameters:
  • A (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

clip(x, a_min, a_max)#

Clip x into [a_min, a_max] (delegates to xp.clip).

Parameters:
  • x (DenseArray)

  • a_min (ArrayLike)

  • a_max (ArrayLike)

Return type:

DenseArray

column_stack(arrays)#

Stack 1-D arrays as columns into a 2-D array (delegates to xp.column_stack).

1-D inputs become columns; higher-dimensional inputs are concatenated along axis 1, matching NumPy column_stack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

concatenate(arrays, axis=0, dtype=None)#

Concatenate arrays along an existing axis (delegates to xp.concat).

Parameters:
  • arrays (Sequence[DenseArray])

  • axis (int)

  • dtype (Any | None)

Return type:

DenseArray

conj(x)#

Complex conjugate of x (delegates to xp.conj).

Parameters:

x (DenseArray)

Return type:

DenseArray

diag(x)#

Extract or construct a diagonal (delegates to xp.diag).

Parameters:

x (DenseArray)

Return type:

DenseArray

diagonal(x)#

Return the main diagonal of x (delegates to xp.diagonal).

Parameters:

x (DenseArray)

Return type:

DenseArray

dstack(arrays)#

Stack arrays depth-wise along the third axis (delegates to xp.dstack).

Inputs are promoted to at least 3-D and concatenated along axis 2, matching NumPy dstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

property e: DenseArray#

Euler’s number as a cached backend scalar.

eigh(x, backend_kwargs=None)#

Eigenpairs of a Hermitian dense matrix (delegates to xp.linalg.eigh).

Parameters:
  • x (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseArray, DenseArray]

eigvalsh(A, backend_kwargs=None)#

Eigenvalues of a Hermitian dense matrix (delegates to xp.linalg.eigvalsh).

Parameters:
  • A (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

einsum(subscripts, *operands)#

Einstein summation (delegates to xp.einsum).

Parameters:
  • subscripts (str)

  • operands (DenseArray)

Return type:

DenseArray

empty(shape, dtype=None)#

Create an uninitialized array (delegates to xp.empty).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

eps(dtype)#

Machine epsilon for dtype.

Parameters:

dtype (Any)

Return type:

float

exp(x)#

Elementwise exponential (delegates to xp.exp).

Parameters:

x (DenseArray)

Return type:

DenseArray

expand_dims(x, axis)#

Insert singleton dimensions into x.

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int])

Return type:

DenseArray

eye(n, m=None, dtype=None)#

Create an identity-like matrix (delegates to xp.eye).

Parameters:
  • n (int)

  • m (int | None)

  • dtype (Any | None)

Return type:

DenseArray

property family: str#

Backend family identifier.

free_memory_bytes()#

Return currently free device memory in bytes, or None if unknown.

The kernel dispatcher (ADR-016) uses this to gate materializing fast paths — those that allocate more than O(1) extra memory — against a memory budget before selecting them. The base implementation returns None (unknown): a backend that can cheaply query free memory (e.g. a GPU runtime) overrides this. When the budget is unknown the dispatcher treats any cost-carrying spec as unaffordable, so reporting None is always safe and never routes to a materializing kernel.

Return type:

int | None

full(shape, fill_value, dtype=None)#

Create a value-filled array (delegates to xp.full).

Parameters:
  • shape (Tuple[int, ...])

  • fill_value (Any)

  • dtype (Any | None)

Return type:

DenseArray

full_like(x, value, dtype=None)#

Create a value-filled array like x (delegates to xp.full_like).

Parameters:
  • x (DenseArray)

  • value (Any)

  • dtype (Any | None)

Return type:

DenseArray

get_dtype(x)#

Return x.dtype after verifying x is a backend array.

Parameters:

x (Any)

Return type:

Any

hstack(arrays)#

Stack arrays horizontally / column-wise (delegates to xp.hstack).

Concatenates along axis 0 for 1-D inputs and along axis 1 otherwise, matching NumPy hstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

imag(x)#

Imaginary component of x (delegates to xp.imag).

Parameters:

x (DenseArray)

Return type:

DenseArray

property inf: DenseArray#

Positive infinity as a cached backend scalar.

is_array(x)#

Return whether x is any array for this backend.

Parameters:

x (Any)

Return type:

bool

is_complex_dtype(dtype)#

Return whether dtype is a complex floating type.

Parameters:

dtype (Any) – Backend or portable dtype specifier to inspect.

Returns:

True when dtype represents complex floating values.

Return type:

bool

is_dense(x)#

Return whether x is a dense array for this backend.

Parameters:

x (Any)

Return type:

bool

is_sparse(x)#

Return whether x is a sparse array for this backend.

Parameters:

x (Any)

Return type:

bool

isfinite(x)#

Elementwise finite check (delegates to xp.isfinite).

Parameters:

x (DenseArray)

Return type:

DenseArray

isnan(x)#

Elementwise NaN check (delegates to xp.isnan).

Parameters:

x (DenseArray)

Return type:

DenseArray

kron(a, b)#

Kronecker product (delegates to xp.kron).

Parameters:
  • a (DenseArray)

  • b (DenseArray)

Return type:

DenseArray

log(x)#

Elementwise natural logarithm (delegates to xp.log).

Parameters:

x (DenseArray)

Return type:

DenseArray

matmul(a, b, backend_kwargs=None)#

Matrix product (delegates to xp.matmul).

Parameters:
  • a (DenseArray)

  • b (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

max(x, axis=None, keepdims=False)#

Maximum over given axes (delegates to xp.max).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

maximum(x, y)#

Elementwise maximum (delegates to xp.maximum).

Parameters:
  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

mean(x, axis=None, keepdims=False)#

Mean over given axes (delegates to xp.mean).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

min(x, axis=None, keepdims=False)#

Minimum over given axes (delegates to xp.min).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

minimum(x, y)#

Elementwise minimum (delegates to xp.minimum).

Parameters:
  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

moveaxis(x, source, destination)#

Move axes of x to new positions.

Parameters:
  • x (DenseArray)

  • source (int | Sequence[int])

  • destination (int | Sequence[int])

Return type:

DenseArray

property nan: DenseArray#

NaN as a cached backend scalar.

ndim(x)#

Return the number of dimensions of x.

Parameters:

x (Any)

Return type:

int

norm(x, ord=None, axis=None, keepdims=False)#

Vector or matrix norm (delegates to xp.linalg.norm).

Parameters:
  • x (DenseArray)

  • ord (int | str | None)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

ones(shape, dtype=None)#

Create a one-filled array (delegates to xp.ones).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

ones_like(x, dtype=None)#

Create a one-filled array like x (delegates to xp.ones_like).

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

Return type:

DenseArray

property pi: DenseArray#

Pi as a cached backend scalar.

prod(x, axis=None, keepdims=False, dtype=None)#

Product over given axes (delegates to xp.prod).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

Return type:

DenseArray

ravel(x)#

Flatten x to one dimension.

Parameters:

x (DenseArray)

Return type:

DenseArray

real(x)#

Real component of x (delegates to xp.real).

Parameters:

x (DenseArray)

Return type:

DenseArray

real_dtype(dtype)#

Return the real floating dtype with the same precision as dtype.

Parameters:

dtype (Any) – Backend or portable dtype specifier.

Returns:

dtype itself when it is already real-valued; otherwise float32 for complex64 and float64 for complex128.

Return type:

DType

reshape(x, shape)#

Reshape x (delegates to xp.reshape).

Parameters:
  • x (DenseArray)

  • shape (Tuple[int, ...] | int)

Return type:

DenseArray

shape(x)#

Return x.shape as a tuple.

Parameters:

x (Any)

Return type:

tuple[int, …]

sign(x)#

Elementwise sign of x (delegates to xp.sign).

Parameters:

x (DenseArray)

Return type:

DenseArray

size(x)#

Return the total number of elements in x.

Parameters:

x (Any)

Return type:

int

solve(A, b, backend_kwargs=None)#

Solve a dense linear system (delegates to xp.linalg.solve).

Parameters:
  • A (DenseArray)

  • b (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

sort(x, axis=-1)#

Sort x along an axis (delegates to xp.sort).

Parameters:
  • x (DenseArray)

  • axis (int)

Return type:

DenseArray

sqrt(x)#

Elementwise square root of x (delegates to xp.sqrt).

Parameters:

x (DenseArray)

Return type:

DenseArray

squeeze(x, axis=None)#

Remove singleton dimensions from x.

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

Return type:

DenseArray

stack(arrays, axis=0)#

Stack arrays along a new axis (delegates to xp.stack).

Parameters:
  • arrays (Sequence[DenseArray])

  • axis (int)

Return type:

DenseArray

sum(x, axis=None, keepdims=False, dtype=None)#

Sum over given axes (delegates to xp.sum).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

Return type:

DenseArray

svd(A, full_matrices=True, backend_kwargs=None)#

Singular value decomposition (delegates to xp.linalg.svd).

Parameters:
  • A (DenseArray)

  • full_matrices (bool)

  • backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseArray, DenseArray, DenseArray]

swapaxes(x, axis1, axis2)#

Swap two axes of x.

Parameters:
  • x (DenseArray)

  • axis1 (int)

  • axis2 (int)

Return type:

DenseArray

take(x, indices, axis=None)#

Take entries from x by integer indices (delegates to xp.take).

Parameters:
  • x (DenseArray)

  • indices (DenseArray)

  • axis (int | None)

Return type:

DenseArray

trace(x)#

Trace of a matrix (delegates to xp.trace when available).

Parameters:

x (DenseArray)

Return type:

DenseArray

transpose(x, axes=None)#

Permute dimensions of x.

Parameters:
  • x (DenseArray)

  • axes (Sequence[int] | None)

Return type:

DenseArray

tril(x)#

Lower triangle of x (delegates to xp.tril).

Parameters:

x (DenseArray)

Return type:

DenseArray

triu(x)#

Upper triangle of x (delegates to xp.triu).

Parameters:

x (DenseArray)

Return type:

DenseArray

vdot(x, y)#

Return sum(conj(x) * y) over flattened inputs.

Matches NumPy, JAX, and Torch vdot semantics. DenseLinOp.rapply relies on this convention for complex inputs.

Parameters:
  • x (DenseArray)

  • y (DenseArray)

Return type:

DenseArray

vectorize(pyfunc, *, excluded=None, signature=None)#

Vectorize a scalar Python function over its array arguments.

Returns a callable that applies pyfunc elementwise, broadcasting the array arguments against one another and preserving the broadcast shape. Mirrors numpy.vectorize().

Parameters:
  • pyfunc (Callable) – Function called on scalar elements of the (broadcast) inputs.

  • excluded (Sequence[int] | None) – Positional argument indices passed through to pyfunc unchanged instead of being vectorized.

  • signature (str | None) – Generalized-ufunc signature (e.g. "(n),(n)->()"). Supported only on backends that provide a native vectorize.

Return type:

Callable

Notes

Delegates to the backend’s native vectorize (NumPy, JAX, CuPy) when available; otherwise applies a portable Python-loop fallback. The fallback does not support signature.

vstack(arrays)#

Stack arrays vertically / row-wise (delegates to xp.vstack).

Inputs are promoted to at least 2-D and concatenated along axis 0, matching NumPy vstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

where(condition, x, y)#

Select between x and y by condition (delegates to xp.where).

Parameters:
  • condition (DenseArray | bool)

  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

zeros(shape, dtype=None)#

Create a zero-filled array (delegates to xp.zeros).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

zeros_like(x, dtype=None)#

Create a zero-filled array like x (delegates to xp.zeros_like).

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

Return type:

DenseArray

class spacecore.backend.TorchOps[source]#

Bases: EagerControlFlowMixin, BackendOps

BackendOps implementation for PyTorch tensors.

This backend uses PyTorch for dense and sparse tensor operations.

Dense arrays

torch.Tensor with strided layout

Sparse arrays

torch.Tensor with a PyTorch sparse layout

Most methods mirror the corresponding PyTorch public API signatures and
delegate to ``torch`` or ``torch.linalg``. Backend-specific behavior,
dtype promotion, broadcasting, device placement, autograd tracking, and
error modes therefore follow PyTorch semantics.
Backend handles
  • torchmodule

    PyTorch module stored on the class and available through instances as ops.torch. Advanced users may use it when SpaceCore’s portable API does not expose a required PyTorch feature.

Notes

Code intended to remain backend-portable should prefer BackendOps methods. Direct use of ops.torch is an explicit PyTorch-specific escape hatch.

TorchOps follows PyTorch dtype defaults. When no dtype is provided, sanitize_dtype(None) returns torch.get_default_dtype(). Python complex maps to torch.complex64 or torch.complex128 based on the active default floating dtype, and NumPy dtype specifiers are mapped to their corresponding PyTorch dtypes when supported.

Array creation and conversion methods may accept a backend-specific device= keyword. Existing tensors stay on their device unless an explicit device conversion is requested. Dense conversion and ordinary math operations do not detach tensors; autograd metadata is preserved according to normal PyTorch rules.

property dense_array: Type[Any]#

Dense array type using PyTorch.

Return type:

Concrete dense tensor class accepted by this backend.

See:

https://docs.pytorch.org/docs/stable/tensors.html

property sparse_array: Tuple[Type[Any], ...]#

Sparse array type tuple using PyTorch.

Return type:

Tensor class accepted by this backend for sparse tensor layouts.

See:

https://docs.pytorch.org/docs/stable/sparse.html

is_dense(x)[source]#

Check whether an object is a dense PyTorch tensor.

Input:

x: Object to inspect.

Output:

Boolean indicating whether x is a strided PyTorch tensor.

See:

https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch-layout

Parameters:

x (Any)

Return type:

bool

is_sparse(x)[source]#

Check whether an object is a sparse PyTorch tensor.

Input:

x: Object to inspect.

Output:

Boolean indicating whether x is a PyTorch tensor with a sparse layout.

See:

https://docs.pytorch.org/docs/stable/sparse.html

Parameters:

x (Any)

Return type:

bool

sanitize_dtype(dtype)[source]#

Normalize a dtype specifier using PyTorch.

Input:

dtype: Optional dtype requested by SpaceCore or the caller.

Output:

Backend dtype object accepted by PyTorch tensor constructors.

See:

https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch-dtype

Backend-specific notes:

None follows torch.get_default_dtype(). NumPy dtype specifiers are mapped to equivalent PyTorch dtypes when supported.

Parameters:

dtype (Any | None)

Return type:

Any

assparse(x, *, format='coo', dtype=None, device=None)[source]#

Convert input to a sparse tensor using PyTorch.

Input:

x: Dense, sparse, or SciPy sparse input plus sparse format, dtype, and device.

Output:

Sparse backend tensor in COO, CSR, or CSC format.

See:

https://docs.pytorch.org/docs/stable/sparse.html

Backend-specific notes:

SciPy sparse inputs are converted through COO indices and values. Dense inputs are converted through PyTorch’s sparse COO conversion.

Parameters:
  • x (Any)

  • format (Literal['coo', 'csr', 'csc'])

  • dtype (Any | None)

  • device (Any | None)

Return type:

SparseArray

asarray(x, dtype=None, *, device=None, copy=None, backend_kwargs=None, **extra_kwargs)[source]#

Convert input to a dense backend array (delegates to xp.asarray).

Parameters:
  • x (Any)

  • dtype (Any | None)

  • device (Any | None)

  • copy (bool | None)

  • backend_kwargs (dict[str, Any] | None)

  • extra_kwargs (Any)

Return type:

DenseArray

astype(x, dtype, *, copy=True, non_blocking=False, memory_format=None, backend_kwargs=None, **extra_kwargs)[source]#

Cast x to dtype, returning x unchanged when dtype is None.

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

  • copy (bool)

  • non_blocking (bool)

  • memory_format (Any | None)

  • backend_kwargs (dict[str, Any] | None)

  • extra_kwargs (Any)

Return type:

DenseArray

empty(shape, dtype=None, *, out=None, layout=None, device=None, requires_grad=False, pin_memory=False, memory_format=None)[source]#

Create an uninitialized array (delegates to xp.empty).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

  • out (DenseArray | None)

  • layout (Any | None)

  • device (Any | None)

  • requires_grad (bool)

  • pin_memory (bool)

  • memory_format (Any | None)

Return type:

DenseArray

zeros(shape, dtype=None, *, out=None, layout=None, device=None, requires_grad=False)[source]#

Create a zero-filled array (delegates to xp.zeros).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

  • out (DenseArray | None)

  • layout (Any | None)

  • device (Any | None)

  • requires_grad (bool)

Return type:

DenseArray

zeros_like(x, dtype=None, *, layout=None, device=None, requires_grad=False, memory_format=None)[source]#

Create a zero-filled array like x (delegates to xp.zeros_like).

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

  • layout (Any | None)

  • device (Any | None)

  • requires_grad (bool)

  • memory_format (Any | None)

Return type:

DenseArray

arange(start, stop=None, step=None, dtype=None, *, out=None, layout=None, device=None, requires_grad=False)[source]#

Create an evenly spaced range (delegates to xp.arange).

Parameters:
  • start (int)

  • stop (int | None)

  • step (int | None)

  • dtype (Any | None)

  • out (DenseArray | None)

  • layout (Any | None)

  • device (Any | None)

  • requires_grad (bool)

Return type:

DenseArray

sum(x, axis=None, keepdims=False, dtype=None, *, out=None)[source]#

Sum over given axes (delegates to xp.sum).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

  • out (DenseArray | None)

Return type:

DenseArray

matmul(a, b, backend_kwargs=None, *, out=None)[source]#

Matrix product (delegates to xp.matmul).

Parameters:
  • a (DenseArray)

  • b (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

  • out (DenseArray | None)

Return type:

DenseArray

sparse_matmul(a, b, *, reduce='sum')[source]#

Matrix-multiply a sparse tensor by a dense tensor using PyTorch.

Input:

a: Sparse backend tensor; b: Dense backend tensor or vector.

Output:

Dense backend tensor.

See:

https://docs.pytorch.org/docs/stable/generated/torch.sparse.mm.html

Parameters:
  • a (SparseArray)

  • b (DenseArray)

  • reduce (Literal['sum', 'mean', 'amax', 'amin'])

Return type:

DenseArray

vmap(fn, in_axes=0, out_axes=0)[source]#

Vectorize a function using PyTorch’s native vmap when available.

Parameters:
  • fn (Callable)

  • in_axes (int | Sequence[int | None] | None)

  • out_axes (int | Sequence[int | None] | None)

Return type:

Callable

property has_native_vmap: bool#

Return True because supported PyTorch versions provide native vmap.

eigh(x, backend_kwargs=None, UPLO='L', *, out=None)[source]#

Eigenpairs of a Hermitian dense matrix (delegates to xp.linalg.eigh).

Parameters:
  • x (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

  • UPLO (Literal['L', 'U'])

  • out (tuple[DenseArray, DenseArray] | None)

Return type:

tuple[DenseArray, DenseArray]

norm(x, ord=None, axis=None, keepdims=False, *, dtype=None, out=None)[source]#

Vector or matrix norm (delegates to xp.linalg.norm).

Parameters:
  • x (DenseArray)

  • ord (int | str | None)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

  • out (DenseArray | None)

Return type:

DenseArray

solve(A, b, backend_kwargs=None, *, left=True, out=None)[source]#

Solve a dense linear system (delegates to xp.linalg.solve).

Parameters:
  • A (DenseArray)

  • b (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

  • left (bool)

  • out (DenseArray | None)

Return type:

DenseArray

svd(A, full_matrices=True, backend_kwargs=None, *, driver=None, out=None)[source]#

Singular value decomposition (delegates to xp.linalg.svd).

Parameters:
  • A (DenseArray)

  • full_matrices (bool)

  • backend_kwargs (dict[str, Any] | None)

  • driver (str | None)

  • out (DenseArray | tuple[DenseArray, DenseArray, DenseArray] | None)

Return type:

tuple[DenseArray, DenseArray, DenseArray]

cholesky(A, backend_kwargs=None, *, upper=False, out=None)[source]#

Cholesky factorization (delegates to xp.linalg.cholesky).

Parameters:
  • A (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

  • upper (bool)

  • out (DenseArray | None)

Return type:

DenseArray

logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False, *, out=None)[source]#

Compute log-sum-exp using PyTorch.

Input:

a: Dense backend tensor; axis, b, keepdims, return_sign: Reduction controls.

Output:

Dense backend tensor, or (value, sign) when return_sign is true.

See:

https://docs.pytorch.org/docs/stable/generated/torch.logsumexp.html

Backend-specific notes:

Weighted and signed variants are implemented in SpaceCore because PyTorch’s public logsumexp does not expose SciPy-style b or return_sign parameters.

Parameters:
  • a (DenseArray)

  • axis (int | Sequence[int] | None)

  • b (DenseArray | None)

  • keepdims (bool)

  • return_sign (bool)

  • out (DenseArray | None)

Return type:

DenseArray | tuple[DenseArray, DenseArray]

where(condition, x, y, *, out=None)[source]#

Select between x and y by condition (delegates to xp.where).

Parameters:
  • condition (DenseArray | bool)

  • x (ArrayLike)

  • y (ArrayLike)

  • out (DenseArray | None)

Return type:

DenseArray

concatenate(arrays, axis=0, dtype=None, *, out=None)[source]#

Concatenate arrays along an existing axis (delegates to xp.concat).

Parameters:
  • arrays (Sequence[DenseArray])

  • axis (int)

  • dtype (Any | None)

  • out (DenseArray | None)

Return type:

DenseArray

ix_(*args)[source]#

Construct open mesh indices using PyTorch.

Input:

args: One-dimensional index arrays or array-like objects.

Output:

Tuple of broadcastable index tensors.

See:

https://docs.pytorch.org/docs/stable/generated/torch.meshgrid.html

Parameters:

args (Any)

Return type:

Any

allclose_sparse(a, b, rtol=1e-05, atol=1e-08)[source]#

Compare sparse tensors elementwise within tolerances using PyTorch.

Input:

a, b: Sparse backend tensors; rtol and atol: Comparison controls.

Output:

Boolean indicating whether sparse tensors are close.

See:

https://docs.pytorch.org/docs/stable/sparse.html

Backend-specific notes:

Sparse tensors are compared by converting both operands to dense tensors before calling allclose.

Parameters:
  • a (SparseArray)

  • b (SparseArray)

  • rtol (float)

  • atol (float)

Return type:

bool

abs(x)#

Absolute value of x (delegates to xp.abs).

Parameters:

x (DenseArray)

Return type:

DenseArray

allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)#

Return whether dense arrays are close within tolerances.

Parameters:
  • a (DenseArray)

  • b (DenseArray)

  • rtol (float)

  • atol (float)

  • equal_nan (bool)

Return type:

bool

property allow_sparse: bool#

Whether this backend supports sparse arrays.

argmax(x, axis=None, keepdims=False)#

Return indices of maxima along an axis.

Parameters:
  • x (DenseArray)

  • axis (int | None)

  • keepdims (bool)

Return type:

DenseArray

argmin(x, axis=None, keepdims=False)#

Return indices of minima along an axis.

Parameters:
  • x (DenseArray)

  • axis (int | None)

  • keepdims (bool)

Return type:

DenseArray

argsort(x, axis=-1)#

Return indices that sort x along an axis.

Parameters:
  • x (DenseArray)

  • axis (int)

Return type:

DenseArray

broadcast_to(x, shape)#

Broadcast x to shape (delegates to xp.broadcast_to).

Parameters:
  • x (DenseArray)

  • shape (Tuple[int, ...])

Return type:

DenseArray

clip(x, a_min, a_max)#

Clip x into [a_min, a_max] (delegates to xp.clip).

Parameters:
  • x (DenseArray)

  • a_min (ArrayLike)

  • a_max (ArrayLike)

Return type:

DenseArray

column_stack(arrays)#

Stack 1-D arrays as columns into a 2-D array (delegates to xp.column_stack).

1-D inputs become columns; higher-dimensional inputs are concatenated along axis 1, matching NumPy column_stack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

cond(pred, true_fun, false_fun, *operands)#

Select a branch eagerly using Python truth-value conversion.

Parameters:
  • pred (bool)

  • true_fun (Callable[[T], R])

  • false_fun (Callable[[T], R])

  • operands (Any)

Return type:

R

conj(x)#

Complex conjugate of x (delegates to xp.conj).

Parameters:

x (DenseArray)

Return type:

DenseArray

diag(x)#

Extract or construct a diagonal (delegates to xp.diag).

Parameters:

x (DenseArray)

Return type:

DenseArray

diagonal(x)#

Return the main diagonal of x (delegates to xp.diagonal).

Parameters:

x (DenseArray)

Return type:

DenseArray

dstack(arrays)#

Stack arrays depth-wise along the third axis (delegates to xp.dstack).

Inputs are promoted to at least 3-D and concatenated along axis 2, matching NumPy dstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

property e: DenseArray#

Euler’s number as a cached backend scalar.

eigvalsh(A, backend_kwargs=None)#

Eigenvalues of a Hermitian dense matrix (delegates to xp.linalg.eigvalsh).

Parameters:
  • A (DenseArray)

  • backend_kwargs (dict[str, Any] | None)

Return type:

DenseArray

einsum(subscripts, *operands)#

Einstein summation (delegates to xp.einsum).

Parameters:
  • subscripts (str)

  • operands (DenseArray)

Return type:

DenseArray

eps(dtype)#

Machine epsilon for dtype.

Parameters:

dtype (Any)

Return type:

float

exp(x)#

Elementwise exponential (delegates to xp.exp).

Parameters:

x (DenseArray)

Return type:

DenseArray

expand_dims(x, axis)#

Insert singleton dimensions into x.

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int])

Return type:

DenseArray

eye(n, m=None, dtype=None)#

Create an identity-like matrix (delegates to xp.eye).

Parameters:
  • n (int)

  • m (int | None)

  • dtype (Any | None)

Return type:

DenseArray

property family: str#

Backend family identifier.

fori_loop(lower, upper, body_fun, init_val)#

Run a counted loop eagerly as a Python for loop.

This executes without JAX tracing semantics; unroll-style parameters are not accepted because eager backends gain nothing from them.

Parameters:
  • lower (int)

  • upper (int)

  • body_fun (Callable[[int, T], T])

  • init_val (T)

Return type:

T

free_memory_bytes()#

Return currently free device memory in bytes, or None if unknown.

The kernel dispatcher (ADR-016) uses this to gate materializing fast paths — those that allocate more than O(1) extra memory — against a memory budget before selecting them. The base implementation returns None (unknown): a backend that can cheaply query free memory (e.g. a GPU runtime) overrides this. When the budget is unknown the dispatcher treats any cost-carrying spec as unaffordable, so reporting None is always safe and never routes to a materializing kernel.

Return type:

int | None

full(shape, fill_value, dtype=None)#

Create a value-filled array (delegates to xp.full).

Parameters:
  • shape (Tuple[int, ...])

  • fill_value (Any)

  • dtype (Any | None)

Return type:

DenseArray

full_like(x, value, dtype=None)#

Create a value-filled array like x (delegates to xp.full_like).

Parameters:
  • x (DenseArray)

  • value (Any)

  • dtype (Any | None)

Return type:

DenseArray

get_dtype(x)#

Return x.dtype after verifying x is a backend array.

Parameters:

x (Any)

Return type:

Any

hstack(arrays)#

Stack arrays horizontally / column-wise (delegates to xp.hstack).

Concatenates along axis 0 for 1-D inputs and along axis 1 otherwise, matching NumPy hstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

imag(x)#

Imaginary component of x (delegates to xp.imag).

Parameters:

x (DenseArray)

Return type:

DenseArray

index_add(x, index, values, *, copy=True)#

Return x with values accumulated into x[index].

With copy=True a mutable copy is updated and returned; with copy=False x is mutated in place. Immutable backends override this method.

Parameters:
  • x (DenseArray)

  • index (int | slice | Any | Tuple[Any, ...])

  • values (DenseArray)

  • copy (bool)

Return type:

DenseArray

index_set(x, index, values, *, copy=True)#

Return x with x[index] set to values.

With copy=True a mutable copy is updated and returned; with copy=False x is mutated in place. Immutable backends override this method.

Parameters:
  • x (DenseArray)

  • index (int | slice | Any | Tuple[Any, ...])

  • values (ArrayLike)

  • copy (bool)

Return type:

DenseArray

property inf: DenseArray#

Positive infinity as a cached backend scalar.

is_array(x)#

Return whether x is any array for this backend.

Parameters:

x (Any)

Return type:

bool

is_complex_dtype(dtype)#

Return whether dtype is a complex floating type.

Parameters:

dtype (Any) – Backend or portable dtype specifier to inspect.

Returns:

True when dtype represents complex floating values.

Return type:

bool

isfinite(x)#

Elementwise finite check (delegates to xp.isfinite).

Parameters:

x (DenseArray)

Return type:

DenseArray

isnan(x)#

Elementwise NaN check (delegates to xp.isnan).

Parameters:

x (DenseArray)

Return type:

DenseArray

kron(a, b)#

Kronecker product (delegates to xp.kron).

Parameters:
  • a (DenseArray)

  • b (DenseArray)

Return type:

DenseArray

log(x)#

Elementwise natural logarithm (delegates to xp.log).

Parameters:

x (DenseArray)

Return type:

DenseArray

max(x, axis=None, keepdims=False)#

Maximum over given axes (delegates to xp.max).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

maximum(x, y)#

Elementwise maximum (delegates to xp.maximum).

Parameters:
  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

mean(x, axis=None, keepdims=False)#

Mean over given axes (delegates to xp.mean).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

min(x, axis=None, keepdims=False)#

Minimum over given axes (delegates to xp.min).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

Return type:

DenseArray

minimum(x, y)#

Elementwise minimum (delegates to xp.minimum).

Parameters:
  • x (ArrayLike)

  • y (ArrayLike)

Return type:

DenseArray

moveaxis(x, source, destination)#

Move axes of x to new positions.

Parameters:
  • x (DenseArray)

  • source (int | Sequence[int])

  • destination (int | Sequence[int])

Return type:

DenseArray

property nan: DenseArray#

NaN as a cached backend scalar.

ndim(x)#

Return the number of dimensions of x.

Parameters:

x (Any)

Return type:

int

ones(shape, dtype=None)#

Create a one-filled array (delegates to xp.ones).

Parameters:
  • shape (Tuple[int, ...])

  • dtype (Any | None)

Return type:

DenseArray

ones_like(x, dtype=None)#

Create a one-filled array like x (delegates to xp.ones_like).

Parameters:
  • x (DenseArray)

  • dtype (Any | None)

Return type:

DenseArray

property pi: DenseArray#

Pi as a cached backend scalar.

prod(x, axis=None, keepdims=False, dtype=None)#

Product over given axes (delegates to xp.prod).

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

  • keepdims (bool)

  • dtype (Any | None)

Return type:

DenseArray

ravel(x)#

Flatten x to one dimension.

Parameters:

x (DenseArray)

Return type:

DenseArray

real(x)#

Real component of x (delegates to xp.real).

Parameters:

x (DenseArray)

Return type:

DenseArray

real_dtype(dtype)#

Return the real floating dtype with the same precision as dtype.

Parameters:

dtype (Any) – Backend or portable dtype specifier.

Returns:

dtype itself when it is already real-valued; otherwise float32 for complex64 and float64 for complex128.

Return type:

DType

reshape(x, shape)#

Reshape x (delegates to xp.reshape).

Parameters:
  • x (DenseArray)

  • shape (Tuple[int, ...] | int)

Return type:

DenseArray

scan(f, init, xs, length=None, reverse=False, unroll=1)#

Run a scan primitive eagerly, stacking per-step outputs at the end.

unroll is accepted only for API parity with jax.lax.scan. When xs is None an explicit length is required; otherwise the length is inferred from the leading axis of the first leaf of xs.

Parameters:
  • f (Callable[[Carry, X], Tuple[Carry, Y]])

  • init (Carry)

  • xs (X)

  • length (int | None)

  • reverse (bool)

  • unroll (int)

Return type:

Tuple[Carry, Y]

shape(x)#

Return x.shape as a tuple.

Parameters:

x (Any)

Return type:

tuple[int, …]

sign(x)#

Elementwise sign of x (delegates to xp.sign).

Parameters:

x (DenseArray)

Return type:

DenseArray

size(x)#

Return the total number of elements in x.

Parameters:

x (Any)

Return type:

int

sort(x, axis=-1)#

Sort x along an axis (delegates to xp.sort).

Parameters:
  • x (DenseArray)

  • axis (int)

Return type:

DenseArray

sqrt(x)#

Elementwise square root of x (delegates to xp.sqrt).

Parameters:

x (DenseArray)

Return type:

DenseArray

squeeze(x, axis=None)#

Remove singleton dimensions from x.

Parameters:
  • x (DenseArray)

  • axis (int | Sequence[int] | None)

Return type:

DenseArray

stack(arrays, axis=0)#

Stack arrays along a new axis (delegates to xp.stack).

Parameters:
  • arrays (Sequence[DenseArray])

  • axis (int)

Return type:

DenseArray

swapaxes(x, axis1, axis2)#

Swap two axes of x.

Parameters:
  • x (DenseArray)

  • axis1 (int)

  • axis2 (int)

Return type:

DenseArray

take(x, indices, axis=None)#

Take entries from x by integer indices (delegates to xp.take).

Parameters:
  • x (DenseArray)

  • indices (DenseArray)

  • axis (int | None)

Return type:

DenseArray

trace(x)#

Trace of a matrix (delegates to xp.trace when available).

Parameters:

x (DenseArray)

Return type:

DenseArray

transpose(x, axes=None)#

Permute dimensions of x.

Parameters:
  • x (DenseArray)

  • axes (Sequence[int] | None)

Return type:

DenseArray

tril(x)#

Lower triangle of x (delegates to xp.tril).

Parameters:

x (DenseArray)

Return type:

DenseArray

triu(x)#

Upper triangle of x (delegates to xp.triu).

Parameters:

x (DenseArray)

Return type:

DenseArray

vdot(x, y)#

Return sum(conj(x) * y) over flattened inputs.

Matches NumPy, JAX, and Torch vdot semantics. DenseLinOp.rapply relies on this convention for complex inputs.

Parameters:
  • x (DenseArray)

  • y (DenseArray)

Return type:

DenseArray

vectorize(pyfunc, *, excluded=None, signature=None)#

Vectorize a scalar Python function over its array arguments.

Returns a callable that applies pyfunc elementwise, broadcasting the array arguments against one another and preserving the broadcast shape. Mirrors numpy.vectorize().

Parameters:
  • pyfunc (Callable) – Function called on scalar elements of the (broadcast) inputs.

  • excluded (Sequence[int] | None) – Positional argument indices passed through to pyfunc unchanged instead of being vectorized.

  • signature (str | None) – Generalized-ufunc signature (e.g. "(n),(n)->()"). Supported only on backends that provide a native vectorize.

Return type:

Callable

Notes

Delegates to the backend’s native vectorize (NumPy, JAX, CuPy) when available; otherwise applies a portable Python-loop fallback. The fallback does not support signature.

vstack(arrays)#

Stack arrays vertically / row-wise (delegates to xp.vstack).

Inputs are promoted to at least 2-D and concatenated along axis 0, matching NumPy vstack semantics.

Parameters:

arrays (Sequence[DenseArray])

Return type:

DenseArray

while_loop(cond_fun, body_fun, init_val)#

Run a while-loop eagerly, converting the predicate to bool each step.

Parameters:
  • cond_fun (Callable[[T], bool])

  • body_fun (Callable[[T], T])

  • init_val (T)

Return type:

T