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`	Backend operations on NumPy arrays, with SciPy sparse matrices and arrays.
`spacecore.backend.JaxOps`	Backend operations on JAX arrays (`jax.Array`; sparse `BCOO` and `BCSR`).
`spacecore.backend.TorchOps`	Backend operations on PyTorch tensors (strided dense, or any sparse layout).

NumpyOps is always available; JaxOps, TorchOps and CuPyOps are exported only when their package is installed.

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.

Array types and editor hints#

Concrete backend operations retain their native dense array type. The same is true for an explicitly constructed context:

ops = sc.NumpyOps()
x = ops.asarray([1.0, 2.0])       # numpy.typing.NDArray[Any]
y = ops.sqrt(x)                  # numpy.typing.NDArray[Any]
ctx = sc.Context(ops)
z = ctx.asarray([3.0, 4.0])      # numpy.typing.NDArray[Any]
z.astype(np.float32)             # native NumPy completions

JaxOps yields jax.Array and TorchOps yields torch.Tensor. Install the corresponding backend in the editor’s Python environment to get its native completions. SpaceCore ships its typing metadata in the wheel.

For portable code, use BackendOps and DenseArray. A heterogeneous list[BackendOps] exposes the common array protocol. A context resolved from a string or ambient state also exposes this common protocol: a registry entry can change at runtime, so its native type cannot be inferred safely from its name. Space and operator APIs continue to expose their portable representations. CuPy and sparse results retain their existing portable protocols.

These annotations track the array family, not exact shape, dtype, device, sparse layout, or gradient state. Runtime space checks still validate those properties. spacecore.ArrayLike means an object with array metadata; unlike NumPy’s ArrayLike, it does not include Python lists. Constructors such as ops.asarray accept lists separately.

Autodoc#

class spacecore.backend.BackendOps[source]#

Bases: ABC, Generic[DenseT]

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.

is_tracing()[source]#

Whether code is running under this backend’s tracer (jit, vmap, compile).

Return type:: bool

classmethod install_pytree_protocol()[source]#

Register SpaceCore containers with this backend’s pytree machinery; a no-op by default.

Backends with transforms (JAX, torch) override this and import their modules inside the method, so a missing dependency cannot break import spacecore.

Return type:: 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:

DenseT

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:

DenseT | Tuple[DenseT, DenseT]

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 | float | complex)
copy (bool)

Return type:

DenseT

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:

DenseT

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: DenseT#: Positive infinity as a cached backend scalar.

property nan: DenseT#: NaN as a cached backend scalar.

property pi: DenseT#: Pi as a cached backend scalar.

property e: DenseT#: 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

complex_dtype(dtype)[source]#

Return the complex dtype of the same precision as dtype (float32 -> complex64).

Needed because result_type(float32, 1j) is complex128 on NumPy but complex64 on JAX and torch.

Parameters:: dtype (Any)
Return type:: Any

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:

DenseT

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:

DenseT

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

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

Parameters:

shape (Tuple[int, ...])
dtype (Any | None)

Return type:

DenseT

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

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

Parameters:

shape (Tuple[int, ...])
dtype (Any | None)

Return type:

DenseT

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

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

Parameters:

shape (Tuple[int, ...])
dtype (Any | None)

Return type:

DenseT

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:

DenseT

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:

DenseT

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:

DenseT

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:

DenseT

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:

DenseT

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:

DenseT

ravel(x)[source]#

Flatten x to one dimension in C order (last axis varies fastest).

Parameters:: x (DenseArray)
Return type:: DenseT

reshape(x, shape)[source]#

Reshape x in C order (delegates to xp.reshape; see the note above).

Parameters:

x (DenseArray)
shape (Tuple[int, ...] | int)

Return type:

DenseT

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

Permute dimensions of x.

Parameters:

x (DenseArray)
axes (Sequence[int] | None)

Return type:

DenseT

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

Swap two axes of x.

Parameters:

x (DenseArray)
axis1 (int)
axis2 (int)

Return type:

DenseT

expand_dims(x, axis)[source]#

Insert singleton dimensions into x.

Parameters:

x (DenseArray)
axis (int | Sequence[int])

Return type:

DenseT

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

Remove singleton dimensions from x.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)

Return type:

DenseT

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:

DenseT

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

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

Parameters:

arrays (Sequence[DenseArray])
axis (int)

Return type:

DenseT

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:: DenseT

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:: DenseT

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:: DenseT

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:: DenseT

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 function over broadcast array arguments, like numpy.vectorize().

Uses the backend’s native vectorize when present; the Python-loop fallback does not support signature.

Parameters:

pyfunc (Callable)
excluded (Sequence[int] | None)
signature (str | None)

Return type:

Callable

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

Sum over given axes (delegates to xp.sum). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)
dtype (Any | None)

Return type:

DenseT

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

Mean over given axes (delegates to xp.mean). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

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

Minimum over given axes (delegates to xp.min). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

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

Maximum over given axes (delegates to xp.max). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

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

Product over given axes (delegates to xp.prod). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)
dtype (Any | None)

Return type:

DenseT

trace(x)[source]#

Trace over the trailing two axes; batched input traces each matrix.

Deliberately not xp.trace(x): NumPy and JAX default to axis1=0, axis2=1, which on a (batch, n, n) stack traces across the batch axis and silently returns a length-n vector, while torch.trace rejects anything but a 2-D input. SpaceCore’s batched convention is leading batch axes, so the trailing two are the matrix.

Parameters:: x (DenseArray)
Return type:: DenseT

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

Return indices that sort x along an axis.

Parameters:

x (DenseArray)
axis (int)

Return type:

DenseT

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

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

Parameters:

x (DenseArray)
axis (int)

Return type:

DenseT

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

Return indices of minima along an axis. Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | None)
keepdims (bool)

Return type:

DenseT

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

Return indices of maxima along an axis. Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | None)
keepdims (bool)

Return type:

DenseT

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. Returns a 0-d array.

Parameters:

x (DenseArray)
y (DenseArray)

Return type:

DenseT

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:

DenseT

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

Eigenpairs of a Hermitian dense matrix, in SpaceCore’s presentation.

Eigenvalues ascending; each eigenvector scaled so its largest-magnitude entry is real and positive. See the conventions note above – this is deliberately not raw xp.linalg.eigh output.

Parameters:

x (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseT, DenseT]

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

Vector or matrix norm (delegates to xp.linalg.norm). Returns a 0-d array when total.

Parameters:

x (DenseArray)
ord (int | str | None)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

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:

DenseT

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

Eigenvalues of a Hermitian dense matrix, ascending (SpaceCore convention).

Parameters:

A (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

DenseT

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

Singular value decomposition, in SpaceCore’s presentation.

Singular values descending; each of the leading k = min(m, n) columns of U scaled so its largest-magnitude entry is real and positive, with the compensating phase applied to the matching row of Vh so that U @ diag(s) @ Vh still reconstructs A exactly.

With full_matrices=True the trailing columns of U (and rows of Vh) span a null space whose basis is arbitrary and backend-specific – they are left untouched, and are not comparable across backends. Pass full_matrices=False when you need agreement.

Parameters:

A (DenseArray)
full_matrices (bool)
backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseT, DenseT, DenseT]

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

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

Parameters:

A (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

DenseT

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

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

Parameters:

condition (DenseArray | bool)
x (ArrayLike | float | complex)
y (ArrayLike | float | complex)

Return type:

DenseT

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:

DenseT

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:

DenseT

diagonal(x)[source]#

Return the main diagonal over the trailing two axes.

Same correction as trace(): the array libraries default to axis1=0, axis2=1, which on a (batch, n, n) stack diagonalizes across the batch axis. Identical to the library default for 2-D input.

Parameters:: x (DenseArray)
Return type:: DenseT

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

conj(x, **kwargs)[source]#

Return xp.conj(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

real(x, **kwargs)[source]#

Return xp.real(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

imag(x, **kwargs)[source]#

Return xp.imag(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

abs(x, **kwargs)[source]#

Return xp.abs(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

sign(x, **kwargs)[source]#

Return xp.sign(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

sqrt(x, **kwargs)[source]#

Return xp.sqrt(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

exp(x, **kwargs)[source]#

Return xp.exp(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

log(x, **kwargs)[source]#

Return xp.log(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

isfinite(x, **kwargs)[source]#

Return xp.isfinite(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

isnan(x, **kwargs)[source]#

Return xp.isnan(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

broadcast_to(x, shape, **kwargs)[source]#

Return xp.broadcast_to(...) as a backend array.

Parameters:

x (DenseArray)
shape (Sequence[int])
kwargs (Any)

Return type:

DenseT

kron(a, b, **kwargs)[source]#

Return xp.kron(...) as a backend array.

Parameters:

a (DenseArray)
b (DenseArray)
kwargs (Any)

Return type:

DenseT

einsum(subscripts, *operands, **kwargs)[source]#

Return xp.einsum(...) as a backend array.

Parameters:

subscripts (str)
operands (DenseArray)
kwargs (Any)

Return type:

DenseT

maximum(x1, x2, **kwargs)[source]#

Return xp.maximum(...) as a backend array.

Parameters:

x1 (DenseArray | float)
x2 (DenseArray | float)
kwargs (Any)

Return type:

DenseT

minimum(x1, x2, **kwargs)[source]#

Return xp.minimum(...) as a backend array.

Parameters:

x1 (DenseArray | float)
x2 (DenseArray | float)
kwargs (Any)

Return type:

DenseT

clip(x, min=None, max=None, **kwargs)[source]#

Return xp.clip(...) as a backend array.

Parameters:

x (DenseArray)
min (Any)
max (Any)
kwargs (Any)

Return type:

DenseT

diag(x, k=0, **kwargs)[source]#

Return xp.diag(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

tril(x, k=0, **kwargs)[source]#

Return xp.tril(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

triu(x, k=0, **kwargs)[source]#

Return xp.triu(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

class spacecore.backend.NumpyOps[source]#

Bases: EagerControlFlowMixin, BackendOps[DenseArray]

Backend operations on NumPy arrays, with SciPy sparse matrices and arrays.

Methods delegate to NumPy and SciPy and follow their semantics; a device argument must be "cpu" or None. ops.np and ops.sp are an explicit escape hatch.

np: Any = <module 'numpy' from '/opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/numpy/__init__.py'>#

sp: Any = <module 'scipy' from '/opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/scipy/__init__.py'>#

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 to a 2-D SciPy sparse matrix in the requested format.

Parameters:

x (Any)
format (Literal['csr', 'csc', 'coo'])
dtype (Any | None)

Return type:

SparseArray

sparse_matmul(a, b)[source]#

Multiply a sparse matrix by a dense array.

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:: Open-mesh indexing documentation

Parameters:: args (Any)
Return type:: Any

abs(x, **kwargs)#

Return xp.abs(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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

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

Compare two SciPy sparse matrices within tolerances, via their CSR difference.

Parameters:

a (SparseArray)
b (SparseArray)
rtol (float)
atol (float)

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:

DenseT

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

Return indices of maxima along an axis. Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | None)
keepdims (bool)

Return type:

DenseT

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

Return indices of minima along an axis. Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | None)
keepdims (bool)

Return type:

DenseT

argsort(x, axis=-1)#

Return indices that sort x along an axis.

Parameters:

x (DenseArray)
axis (int)

Return type:

DenseT

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:

DenseT

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:

DenseT

broadcast_to(x, shape, **kwargs)#

Return xp.broadcast_to(...) as a backend array.

Parameters:

x (DenseArray)
shape (Sequence[int])
kwargs (Any)

Return type:

DenseT

cholesky(A, backend_kwargs=None)#

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

Parameters:

A (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

DenseT

clip(x, min=None, max=None, **kwargs)#

Return xp.clip(...) as a backend array.

Parameters:

x (DenseArray)
min (Any)
max (Any)
kwargs (Any)

Return type:

DenseT

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:: DenseT

complex_dtype(dtype)#

Return the complex dtype of the same precision as dtype (float32 -> complex64).

Needed because result_type(float32, 1j) is complex128 on NumPy but complex64 on JAX and torch.

Parameters:: dtype (Any)
Return type:: Any

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:

DenseT

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, **kwargs)#

Return xp.conj(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

diag(x, k=0, **kwargs)#

Return xp.diag(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

diagonal(x)#

Return the main diagonal over the trailing two axes.

Same correction as trace(): the array libraries default to axis1=0, axis2=1, which on a (batch, n, n) stack diagonalizes across the batch axis. Identical to the library default for 2-D input.

Parameters:: x (DenseArray)
Return type:: DenseT

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:: DenseT

property e: DenseT#: Euler’s number as a cached backend scalar.

eigh(x, backend_kwargs=None)#

Eigenpairs of a Hermitian dense matrix, in SpaceCore’s presentation.

Eigenvalues ascending; each eigenvector scaled so its largest-magnitude entry is real and positive. See the conventions note above – this is deliberately not raw xp.linalg.eigh output.

Parameters:

x (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseT, DenseT]

eigvalsh(A, backend_kwargs=None)#

Eigenvalues of a Hermitian dense matrix, ascending (SpaceCore convention).

Parameters:

A (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

DenseT

einsum(subscripts, *operands, **kwargs)#

Return xp.einsum(...) as a backend array.

Parameters:

subscripts (str)
operands (DenseArray)
kwargs (Any)

Return type:

DenseT

empty(shape, dtype=None)#

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

Parameters:

shape (Tuple[int, ...])
dtype (Any | None)

Return type:

DenseT

eps(dtype)#

Machine epsilon for dtype.

Parameters:: dtype (Any)
Return type:: float

exp(x, **kwargs)#

Return xp.exp(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

expand_dims(x, axis)#

Insert singleton dimensions into x.

Parameters:

x (DenseArray)
axis (int | Sequence[int])

Return type:

DenseT

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:

DenseT

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:

DenseT

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:

DenseT

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:: DenseT

imag(x, **kwargs)#

Return xp.imag(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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:

DenseT

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 | float | complex)
copy (bool)

Return type:

DenseT

property inf: DenseT#: Positive infinity as a cached backend scalar.

classmethod install_pytree_protocol()#

Register SpaceCore containers with this backend’s pytree machinery; a no-op by default.

Backends with transforms (JAX, torch) override this and import their modules inside the method, so a missing dependency cannot break import spacecore.

Return type:: None

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

is_tracing()#

Whether code is running under this backend’s tracer (jit, vmap, compile).

Return type:: bool

isfinite(x, **kwargs)#

Return xp.isfinite(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

isnan(x, **kwargs)#

Return xp.isnan(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

kron(a, b, **kwargs)#

Return xp.kron(...) as a backend array.

Parameters:

a (DenseArray)
b (DenseArray)
kwargs (Any)

Return type:

DenseT

log(x, **kwargs)#

Return xp.log(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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:

DenseT

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

Maximum over given axes (delegates to xp.max). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

maximum(x1, x2, **kwargs)#

Return xp.maximum(...) as a backend array.

Parameters:

x1 (DenseArray | float)
x2 (DenseArray | float)
kwargs (Any)

Return type:

DenseT

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

Mean over given axes (delegates to xp.mean). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

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

Minimum over given axes (delegates to xp.min). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

minimum(x1, x2, **kwargs)#

Return xp.minimum(...) as a backend array.

Parameters:

x1 (DenseArray | float)
x2 (DenseArray | float)
kwargs (Any)

Return type:

DenseT

moveaxis(x, source, destination)#

Move axes of x to new positions.

Parameters:

x (DenseArray)
source (int | Sequence[int])
destination (int | Sequence[int])

Return type:

DenseT

property nan: DenseT#: 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). Returns a 0-d array when total.

Parameters:

x (DenseArray)
ord (int | str | None)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

ones(shape, dtype=None)#

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

Parameters:

shape (Tuple[int, ...])
dtype (Any | None)

Return type:

DenseT

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:

DenseT

property pi: DenseT#: Pi as a cached backend scalar.

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

Product over given axes (delegates to xp.prod). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)
dtype (Any | None)

Return type:

DenseT

ravel(x)#

Flatten x to one dimension in C order (last axis varies fastest).

Parameters:: x (DenseArray)
Return type:: DenseT

real(x, **kwargs)#

Return xp.real(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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 in C order (delegates to xp.reshape; see the note above).

Parameters:

x (DenseArray)
shape (Tuple[int, ...] | int)

Return type:

DenseT

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, **kwargs)#

Return xp.sign(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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:

DenseT

sort(x, axis=-1)#

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

Parameters:

x (DenseArray)
axis (int)

Return type:

DenseT

sqrt(x, **kwargs)#

Return xp.sqrt(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

squeeze(x, axis=None)#

Remove singleton dimensions from x.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)

Return type:

DenseT

stack(arrays, axis=0)#

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

Parameters:

arrays (Sequence[DenseArray])
axis (int)

Return type:

DenseT

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

Sum over given axes (delegates to xp.sum). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)
dtype (Any | None)

Return type:

DenseT

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

Singular value decomposition, in SpaceCore’s presentation.

Singular values descending; each of the leading k = min(m, n) columns of U scaled so its largest-magnitude entry is real and positive, with the compensating phase applied to the matching row of Vh so that U @ diag(s) @ Vh still reconstructs A exactly.

With full_matrices=True the trailing columns of U (and rows of Vh) span a null space whose basis is arbitrary and backend-specific – they are left untouched, and are not comparable across backends. Pass full_matrices=False when you need agreement.

Parameters:

A (DenseArray)
full_matrices (bool)
backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseT, DenseT, DenseT]

swapaxes(x, axis1, axis2)#

Swap two axes of x.

Parameters:

x (DenseArray)
axis1 (int)
axis2 (int)

Return type:

DenseT

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:

DenseT

trace(x)#

Trace over the trailing two axes; batched input traces each matrix.

Deliberately not xp.trace(x): NumPy and JAX default to axis1=0, axis2=1, which on a (batch, n, n) stack traces across the batch axis and silently returns a length-n vector, while torch.trace rejects anything but a 2-D input. SpaceCore’s batched convention is leading batch axes, so the trailing two are the matrix.

Parameters:: x (DenseArray)
Return type:: DenseT

transpose(x, axes=None)#

Permute dimensions of x.

Parameters:

x (DenseArray)
axes (Sequence[int] | None)

Return type:

DenseT

tril(x, k=0, **kwargs)#

Return xp.tril(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

triu(x, k=0, **kwargs)#

Return xp.triu(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

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. Returns a 0-d array.

Parameters:

x (DenseArray)
y (DenseArray)

Return type:

DenseT

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

Vectorize a scalar function over broadcast array arguments, like numpy.vectorize().

Uses the backend’s native vectorize when present; the Python-loop fallback does not support signature.

Parameters:

pyfunc (Callable)
excluded (Sequence[int] | None)
signature (str | None)

Return type:

Callable

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:: DenseT

where(condition, x, y)#

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

Parameters:

condition (DenseArray | bool)
x (ArrayLike | float | complex)
y (ArrayLike | float | complex)

Return type:

DenseT

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:

DenseT

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:

DenseT

class spacecore.backend.JaxOps[source]#

Bases: BackendOps[DenseArray]

Backend operations on JAX arrays (jax.Array; sparse BCOO and BCSR).

Methods delegate to jax.numpy and related modules, so dtypes, tracing, devices and sharding follow JAX. ops.jax, ops.jnp and ops.jsparse expose the modules as an explicit JAX-only escape hatch.

classmethod install_pytree_protocol()[source]#

Register SpaceCore containers as JAX pytree nodes.

JAX’s register_pytree_node_class consumes exactly the protocol PyTreeNode defines — tree_flatten returning (children, aux) and a tree_unflatten(aux, children) classmethod — so it is the registrar and no translation is needed.

Return type:: None

sanitize_dtype(dtype)[source]#

Normalize a dtype for JAX, rejecting dtypes 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 to a JAX sparse array (BCOO/BCSR); SciPy inputs go through 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 a JAX sparse array by a dense one, returning a dense jax.Array.

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

is_tracing()[source]#

Whether a JAX transform (jit, vmap, grad, …) is tracing.

JAX has no public query; trace_state_clean is private, so a JAX that drops it reads as not tracing and checks run as in eager code.

Return type:: bool

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]#

Return x with x[index] = values via x.at[index].set; copy=False raises, as JAX arrays are immutable.

Parameters:

x (DenseArray)
index (int | slice | Any | Tuple[Any, ...])
values (ArrayLike | float | complex)
copy (bool)

Return type:

DenseArray

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:: Open-mesh indexing documentation

Parameters:: args (Any)
Return type:: Any

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

Run jax.lax.fori_loop.

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 jax.lax.while_loop.

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 jax.lax.scan; carries and outputs may be pytrees.

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 jax.lax.cond: both branches are traced, not chosen eagerly in Python.

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]#

Return x with values added at index via x.at[index].add; repeated indices accumulate, and copy=False raises.

Parameters:

x (DenseArray)
index (int | slice | Any | Tuple[Any, ...])
values (DenseArray)
copy (bool)

Return type:

DenseArray

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

Compare two JAX sparse arrays within tolerances, via SciPy sparse.

Parameters:

a (SparseArray)
b (SparseArray)
rtol (float)
atol (float)

Return type:

bool

abs(x, **kwargs)#

Return xp.abs(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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:

DenseT

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

Return indices of maxima along an axis. Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | None)
keepdims (bool)

Return type:

DenseT

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

Return indices of minima along an axis. Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | None)
keepdims (bool)

Return type:

DenseT

argsort(x, axis=-1)#

Return indices that sort x along an axis.

Parameters:

x (DenseArray)
axis (int)

Return type:

DenseT

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:

DenseT

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:

DenseT

broadcast_to(x, shape, **kwargs)#

Return xp.broadcast_to(...) as a backend array.

Parameters:

x (DenseArray)
shape (Sequence[int])
kwargs (Any)

Return type:

DenseT

cholesky(A, backend_kwargs=None)#

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

Parameters:

A (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

DenseT

clip(x, min=None, max=None, **kwargs)#

Return xp.clip(...) as a backend array.

Parameters:

x (DenseArray)
min (Any)
max (Any)
kwargs (Any)

Return type:

DenseT

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:: DenseT

complex_dtype(dtype)#

Return the complex dtype of the same precision as dtype (float32 -> complex64).

Needed because result_type(float32, 1j) is complex128 on NumPy but complex64 on JAX and torch.

Parameters:: dtype (Any)
Return type:: Any

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:

DenseT

conj(x, **kwargs)#

Return xp.conj(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

diag(x, k=0, **kwargs)#

Return xp.diag(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

diagonal(x)#

Return the main diagonal over the trailing two axes.

Same correction as trace(): the array libraries default to axis1=0, axis2=1, which on a (batch, n, n) stack diagonalizes across the batch axis. Identical to the library default for 2-D input.

Parameters:: x (DenseArray)
Return type:: DenseT

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:: DenseT

property e: DenseT#: Euler’s number as a cached backend scalar.

eigh(x, backend_kwargs=None)#

Eigenpairs of a Hermitian dense matrix, in SpaceCore’s presentation.

Eigenvalues ascending; each eigenvector scaled so its largest-magnitude entry is real and positive. See the conventions note above – this is deliberately not raw xp.linalg.eigh output.

Parameters:

x (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseT, DenseT]

eigvalsh(A, backend_kwargs=None)#

Eigenvalues of a Hermitian dense matrix, ascending (SpaceCore convention).

Parameters:

A (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

DenseT

einsum(subscripts, *operands, **kwargs)#

Return xp.einsum(...) as a backend array.

Parameters:

subscripts (str)
operands (DenseArray)
kwargs (Any)

Return type:

DenseT

empty(shape, dtype=None)#

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

Parameters:

shape (Tuple[int, ...])
dtype (Any | None)

Return type:

DenseT

eps(dtype)#

Machine epsilon for dtype.

Parameters:: dtype (Any)
Return type:: float

exp(x, **kwargs)#

Return xp.exp(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

expand_dims(x, axis)#

Insert singleton dimensions into x.

Parameters:

x (DenseArray)
axis (int | Sequence[int])

Return type:

DenseT

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:

DenseT

property family: str#: Backend family identifier.

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:

DenseT

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:

DenseT

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:: DenseT

imag(x, **kwargs)#

Return xp.imag(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

property inf: DenseT#: 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, **kwargs)#

Return xp.isfinite(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

isnan(x, **kwargs)#

Return xp.isnan(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

kron(a, b, **kwargs)#

Return xp.kron(...) as a backend array.

Parameters:

a (DenseArray)
b (DenseArray)
kwargs (Any)

Return type:

DenseT

log(x, **kwargs)#

Return xp.log(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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:

DenseT

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

Maximum over given axes (delegates to xp.max). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

maximum(x1, x2, **kwargs)#

Return xp.maximum(...) as a backend array.

Parameters:

x1 (DenseArray | float)
x2 (DenseArray | float)
kwargs (Any)

Return type:

DenseT

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

Mean over given axes (delegates to xp.mean). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

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

Minimum over given axes (delegates to xp.min). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

minimum(x1, x2, **kwargs)#

Return xp.minimum(...) as a backend array.

Parameters:

x1 (DenseArray | float)
x2 (DenseArray | float)
kwargs (Any)

Return type:

DenseT

moveaxis(x, source, destination)#

Move axes of x to new positions.

Parameters:

x (DenseArray)
source (int | Sequence[int])
destination (int | Sequence[int])

Return type:

DenseT

property nan: DenseT#: 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). Returns a 0-d array when total.

Parameters:

x (DenseArray)
ord (int | str | None)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

ones(shape, dtype=None)#

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

Parameters:

shape (Tuple[int, ...])
dtype (Any | None)

Return type:

DenseT

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:

DenseT

property pi: DenseT#: Pi as a cached backend scalar.

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

Product over given axes (delegates to xp.prod). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)
dtype (Any | None)

Return type:

DenseT

ravel(x)#

Flatten x to one dimension in C order (last axis varies fastest).

Parameters:: x (DenseArray)
Return type:: DenseT

real(x, **kwargs)#

Return xp.real(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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 in C order (delegates to xp.reshape; see the note above).

Parameters:

x (DenseArray)
shape (Tuple[int, ...] | int)

Return type:

DenseT

shape(x)#

Return x.shape as a tuple.

Parameters:: x (Any)
Return type:: tuple[int, …]

sign(x, **kwargs)#

Return xp.sign(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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:

DenseT

sort(x, axis=-1)#

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

Parameters:

x (DenseArray)
axis (int)

Return type:

DenseT

sqrt(x, **kwargs)#

Return xp.sqrt(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

squeeze(x, axis=None)#

Remove singleton dimensions from x.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)

Return type:

DenseT

stack(arrays, axis=0)#

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

Parameters:

arrays (Sequence[DenseArray])
axis (int)

Return type:

DenseT

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

Sum over given axes (delegates to xp.sum). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)
dtype (Any | None)

Return type:

DenseT

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

Singular value decomposition, in SpaceCore’s presentation.

Singular values descending; each of the leading k = min(m, n) columns of U scaled so its largest-magnitude entry is real and positive, with the compensating phase applied to the matching row of Vh so that U @ diag(s) @ Vh still reconstructs A exactly.

With full_matrices=True the trailing columns of U (and rows of Vh) span a null space whose basis is arbitrary and backend-specific – they are left untouched, and are not comparable across backends. Pass full_matrices=False when you need agreement.

Parameters:

A (DenseArray)
full_matrices (bool)
backend_kwargs (dict[str, Any] | None)

Return type:

tuple[DenseT, DenseT, DenseT]

swapaxes(x, axis1, axis2)#

Swap two axes of x.

Parameters:

x (DenseArray)
axis1 (int)
axis2 (int)

Return type:

DenseT

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:

DenseT

trace(x)#

Trace over the trailing two axes; batched input traces each matrix.

Deliberately not xp.trace(x): NumPy and JAX default to axis1=0, axis2=1, which on a (batch, n, n) stack traces across the batch axis and silently returns a length-n vector, while torch.trace rejects anything but a 2-D input. SpaceCore’s batched convention is leading batch axes, so the trailing two are the matrix.

Parameters:: x (DenseArray)
Return type:: DenseT

transpose(x, axes=None)#

Permute dimensions of x.

Parameters:

x (DenseArray)
axes (Sequence[int] | None)

Return type:

DenseT

tril(x, k=0, **kwargs)#

Return xp.tril(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

triu(x, k=0, **kwargs)#

Return xp.triu(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

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. Returns a 0-d array.

Parameters:

x (DenseArray)
y (DenseArray)

Return type:

DenseT

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

Vectorize a scalar function over broadcast array arguments, like numpy.vectorize().

Uses the backend’s native vectorize when present; the Python-loop fallback does not support signature.

Parameters:

pyfunc (Callable)
excluded (Sequence[int] | None)
signature (str | 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:: DenseT

where(condition, x, y)#

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

Parameters:

condition (DenseArray | bool)
x (ArrayLike | float | complex)
y (ArrayLike | float | complex)

Return type:

DenseT

zeros(shape, dtype=None)#

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

Parameters:

shape (Tuple[int, ...])
dtype (Any | None)

Return type:

DenseT

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:

DenseT

class spacecore.backend.TorchOps[source]#

Bases: EagerControlFlowMixin, BackendOps[DenseArray]

Backend operations on PyTorch tensors (strided dense, or any sparse layout).

Methods delegate to torch and torch.linalg, so dtypes, devices and autograd follow PyTorch; sanitize_dtype(None) is torch.get_default_dtype(). ops.torch is an explicit PyTorch-only escape hatch.

classmethod install_pytree_protocol()[source]#

Register SpaceCore containers as torch pytree nodes.

Torch takes children as a list and calls unflatten(children, context), the reverse of PyTreeNode’s order; this adapter translates. torch.utils._pytree is private API. Registering there also covers the C++ pytree, and a breaking change degrades to a registry warning.

Return type:: None

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 for torch; None gives torch.get_default_dtype() and NumPy dtypes map to their torch equivalents.

Parameters:: dtype (Any | None)
Return type:: Any

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

Convert to a torch sparse tensor (COO, CSR or CSC); SciPy inputs go through COO indices and values.

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). Returns a 0-d array when total.

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

is_tracing()[source]#

Whether torch.compile is tracing; Dynamo folds this to a constant in the graph.

Return type:: bool

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 torch.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 supported PyTorch versions provide native vmap.

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

Eigenpairs of a Hermitian dense matrix, in SpaceCore’s presentation.

Eigenvalues ascending; each eigenvector scaled so its largest-magnitude entry is real and positive. See the conventions note above – this is deliberately not raw xp.linalg.eigh output.

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). Returns a 0-d array when total.

Parameters:

x (DenseArray)
ord (int | str | None)
axis (int | Sequence[int] | None)
keepdims (bool)
dtype (Any | None)
out (DenseArray | None)

Return type:

DenseArray

diagonal(x)[source]#

Return the main diagonal over the trailing two axes.

Same correction as trace(): the array libraries default to axis1=0, axis2=1, which on a (batch, n, n) stack diagonalizes across the batch axis. Identical to the library default for 2-D input.

Parameters:: x (DenseArray)
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, in SpaceCore’s presentation.

Singular values descending; each of the leading k = min(m, n) columns of U scaled so its largest-magnitude entry is real and positive, with the compensating phase applied to the matching row of Vh so that U @ diag(s) @ Vh still reconstructs A exactly.

With full_matrices=True the trailing columns of U (and rows of Vh) span a null space whose basis is arbitrary and backend-specific – they are left untouched, and are not comparable across backends. Pass full_matrices=False when you need agreement.

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, with SciPy’s b and return_sign options, which torch lacks, implemented here.

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 | float | complex)
y (ArrayLike | float | complex)
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

abs(x, **kwargs)#

Return xp.abs(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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

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

Compare two torch sparse tensors within tolerances, by densifying both.

Parameters:

a (SparseArray)
b (SparseArray)
rtol (float)
atol (float)

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. Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | None)
keepdims (bool)

Return type:

DenseT

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

Return indices of minima along an axis. Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | None)
keepdims (bool)

Return type:

DenseT

argsort(x, axis=-1)#

Return indices that sort x along an axis.

Parameters:

x (DenseArray)
axis (int)

Return type:

DenseT

broadcast_to(x, shape, **kwargs)#

Return xp.broadcast_to(...) as a backend array.

Parameters:

x (DenseArray)
shape (Sequence[int])
kwargs (Any)

Return type:

DenseT

clip(x, min=None, max=None, **kwargs)#

Return xp.clip(...) as a backend array.

Parameters:

x (DenseArray)
min (Any)
max (Any)
kwargs (Any)

Return type:

DenseT

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:: DenseT

complex_dtype(dtype)#

Return the complex dtype of the same precision as dtype (float32 -> complex64).

Needed because result_type(float32, 1j) is complex128 on NumPy but complex64 on JAX and torch.

Parameters:: dtype (Any)
Return type:: Any

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, **kwargs)#

Return xp.conj(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

diag(x, k=0, **kwargs)#

Return xp.diag(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

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:: DenseT

property e: DenseT#: Euler’s number as a cached backend scalar.

eigvalsh(A, backend_kwargs=None)#

Eigenvalues of a Hermitian dense matrix, ascending (SpaceCore convention).

Parameters:

A (DenseArray)
backend_kwargs (dict[str, Any] | None)

Return type:

DenseT

einsum(subscripts, *operands, **kwargs)#

Return xp.einsum(...) as a backend array.

Parameters:

subscripts (str)
operands (DenseArray)
kwargs (Any)

Return type:

DenseT

eps(dtype)#

Machine epsilon for dtype.

Parameters:: dtype (Any)
Return type:: float

exp(x, **kwargs)#

Return xp.exp(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

expand_dims(x, axis)#

Insert singleton dimensions into x.

Parameters:

x (DenseArray)
axis (int | Sequence[int])

Return type:

DenseT

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:

DenseT

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:

DenseT

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:

DenseT

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:: DenseT

imag(x, **kwargs)#

Return xp.imag(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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:

DenseT

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 | float | complex)
copy (bool)

Return type:

DenseT

property inf: DenseT#: 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, **kwargs)#

Return xp.isfinite(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

isnan(x, **kwargs)#

Return xp.isnan(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

kron(a, b, **kwargs)#

Return xp.kron(...) as a backend array.

Parameters:

a (DenseArray)
b (DenseArray)
kwargs (Any)

Return type:

DenseT

log(x, **kwargs)#

Return xp.log(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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

Maximum over given axes (delegates to xp.max). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

maximum(x1, x2, **kwargs)#

Return xp.maximum(...) as a backend array.

Parameters:

x1 (DenseArray | float)
x2 (DenseArray | float)
kwargs (Any)

Return type:

DenseT

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

Mean over given axes (delegates to xp.mean). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

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

Minimum over given axes (delegates to xp.min). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)

Return type:

DenseT

minimum(x1, x2, **kwargs)#

Return xp.minimum(...) as a backend array.

Parameters:

x1 (DenseArray | float)
x2 (DenseArray | float)
kwargs (Any)

Return type:

DenseT

moveaxis(x, source, destination)#

Move axes of x to new positions.

Parameters:

x (DenseArray)
source (int | Sequence[int])
destination (int | Sequence[int])

Return type:

DenseT

property nan: DenseT#: 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:

DenseT

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:

DenseT

property pi: DenseT#: Pi as a cached backend scalar.

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

Product over given axes (delegates to xp.prod). Returns a 0-d array when total.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)
keepdims (bool)
dtype (Any | None)

Return type:

DenseT

ravel(x)#

Flatten x to one dimension in C order (last axis varies fastest).

Parameters:: x (DenseArray)
Return type:: DenseT

real(x, **kwargs)#

Return xp.real(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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 in C order (delegates to xp.reshape; see the note above).

Parameters:

x (DenseArray)
shape (Tuple[int, ...] | int)

Return type:

DenseT

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, **kwargs)#

Return xp.sign(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

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:

DenseT

sqrt(x, **kwargs)#

Return xp.sqrt(...) as a backend array.

Parameters:

x (DenseArray)
kwargs (Any)

Return type:

DenseT

squeeze(x, axis=None)#

Remove singleton dimensions from x.

Parameters:

x (DenseArray)
axis (int | Sequence[int] | None)

Return type:

DenseT

stack(arrays, axis=0)#

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

Parameters:

arrays (Sequence[DenseArray])
axis (int)

Return type:

DenseT

swapaxes(x, axis1, axis2)#

Swap two axes of x.

Parameters:

x (DenseArray)
axis1 (int)
axis2 (int)

Return type:

DenseT

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:

DenseT

trace(x)#

Trace over the trailing two axes; batched input traces each matrix.

Deliberately not xp.trace(x): NumPy and JAX default to axis1=0, axis2=1, which on a (batch, n, n) stack traces across the batch axis and silently returns a length-n vector, while torch.trace rejects anything but a 2-D input. SpaceCore’s batched convention is leading batch axes, so the trailing two are the matrix.

Parameters:: x (DenseArray)
Return type:: DenseT

transpose(x, axes=None)#

Permute dimensions of x.

Parameters:

x (DenseArray)
axes (Sequence[int] | None)

Return type:

DenseT

tril(x, k=0, **kwargs)#

Return xp.tril(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

triu(x, k=0, **kwargs)#

Return xp.triu(...) as a backend array.

Parameters:

x (DenseArray)
k (int)
kwargs (Any)

Return type:

DenseT

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. Returns a 0-d array.

Parameters:

x (DenseArray)
y (DenseArray)

Return type:

DenseT

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

Vectorize a scalar function over broadcast array arguments, like numpy.vectorize().

Uses the backend’s native vectorize when present; the Python-loop fallback does not support signature.

Parameters:

pyfunc (Callable)
excluded (Sequence[int] | None)
signature (str | 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:: DenseT

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