Backend API#
Backend objects implement the numerical operations used by spaces, operators, functionals, and solvers.
Backend families#
Public numerical contract for SpaceCore backends. |
|
BackendOps implementation for the NumPy/SciPy ecosystem. |
|
BackendOps implementation for the JAX ecosystem. |
|
BackendOps implementation for PyTorch tensors. |
BackendOpsis the abstract operation surface.NumpyOpsis always available and uses NumPy plus SciPy sparse.JaxOpsis optional and follows JAX tracing and dtype configuration.TorchOpsis optional and follows PyTorch tensor, dtype, device, and autograd semantics.CuPyOpsis 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:
ABCPublic numerical contract for SpaceCore backends.
Common dense-array operations delegate to the backend’s Array API-compatible
xpnamespace. 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
vmapis implemented by the backend rather than a Python loop.
- free_memory_bytes()[source]#
Return currently free device memory in bytes, or
Noneif 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 returnsNone(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 reportingNoneis 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
xwithx[index]set tovalues.With
copy=Truea mutable copy is updated and returned; withcopy=Falsexis 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
xwithvaluesaccumulated intox[index].With
copy=Truea mutable copy is updated and returned; withcopy=Falsexis 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.
- is_complex_dtype(dtype)[source]#
Return whether
dtypeis a complex floating type.- Parameters:
dtype (Any) – Backend or portable dtype specifier to inspect.
- Returns:
Truewhendtyperepresents 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:
dtypeitself when it is already real-valued; otherwisefloat32for complex64 andfloat64for complex128.- Return type:
DType
- get_dtype(x)[source]#
Return x.dtype after verifying x is a backend array.
- Parameters:
x (Any)
- Return type:
Any
- 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
- 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
hstacksemantics.- 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
vstacksemantics.- 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
dstacksemantics.- 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_stacksemantics.- Parameters:
arrays (Sequence[DenseArray])
- Return type:
DenseArray
- vmap(fn, in_axes=0, out_axes=0)[source]#
Vectorize
fnover 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
pyfuncelementwise, broadcasting the array arguments against one another and preserving the broadcast shape. Mirrorsnumpy.vectorize().- Parameters:
pyfunc (Callable) – Function called on scalar elements of the (broadcast) inputs.
excluded (Sequence[int] | None) – Positional argument indices passed through to
pyfuncunchanged instead of being vectorized.signature (str | None) – Generalized-ufunc signature (e.g.
"(n),(n)->()"). Supported only on backends that provide a nativevectorize.
- 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 supportsignature.
- 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
xalong 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
vdotsemantics.DenseLinOp.rapplyrelies 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
- class spacecore.backend.NumpyOps[source]#
Bases:
EagerControlFlowMixin,BackendOpsBackendOps 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.
- property sparse_array: Tuple[Type[Any], ...]#
Sparse array type tuple using SciPy.
- Return type:
Concrete sparse array classes accepted by this backend, or None.
- 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:
- 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:
- 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:
- 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
xalong 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_stacksemantics.- 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
dstacksemantics.- 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
forloop.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
vmapis 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
hstacksemantics.- 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
xwithvaluesaccumulated intox[index].With
copy=Truea mutable copy is updated and returned; withcopy=Falsexis 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
xwithx[index]set tovalues.With
copy=Truea mutable copy is updated and returned; withcopy=Falsexis 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
dtypeis a complex floating type.- Parameters:
dtype (Any) – Backend or portable dtype specifier to inspect.
- Returns:
Truewhendtyperepresents 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:
dtypeitself when it is already real-valued; otherwisefloat32for complex64 andfloat64for 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.
unrollis accepted only for API parity withjax.lax.scan. WhenxsisNonean explicitlengthis required; otherwise the length is inferred from the leading axis of the first leaf ofxs.- 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
vdotsemantics.DenseLinOp.rapplyrelies 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
pyfuncelementwise, broadcasting the array arguments against one another and preserving the broadcast shape. Mirrorsnumpy.vectorize().- Parameters:
pyfunc (Callable) – Function called on scalar elements of the (broadcast) inputs.
excluded (Sequence[int] | None) – Positional argument indices passed through to
pyfuncunchanged instead of being vectorized.signature (str | None) – Generalized-ufunc signature (e.g.
"(n),(n)->()"). Supported only on backends that provide a nativevectorize.
- 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 supportsignature.
- vmap(fn, in_axes=0, out_axes=0)#
Vectorize
fnover 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
vstacksemantics.- 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
booleach 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:
BackendOpsBackendOps 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.
- property sparse_array: Tuple[Type[Any], ...]#
Sparse array type tuple using JAX.
- Return type:
Concrete sparse array classes accepted by this backend, or None.
- 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:
- 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:
- 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
Truebecause JAX provides nativevmap.
- 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:
- 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
xalong 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_stacksemantics.- 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
dstacksemantics.- 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
Noneif 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 returnsNone(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 reportingNoneis 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
hstacksemantics.- 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
dtypeis a complex floating type.- Parameters:
dtype (Any) – Backend or portable dtype specifier to inspect.
- Returns:
Truewhendtyperepresents 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:
dtypeitself when it is already real-valued; otherwisefloat32for complex64 andfloat64for 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
vdotsemantics.DenseLinOp.rapplyrelies 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
pyfuncelementwise, broadcasting the array arguments against one another and preserving the broadcast shape. Mirrorsnumpy.vectorize().- Parameters:
pyfunc (Callable) – Function called on scalar elements of the (broadcast) inputs.
excluded (Sequence[int] | None) – Positional argument indices passed through to
pyfuncunchanged instead of being vectorized.signature (str | None) – Generalized-ufunc signature (e.g.
"(n),(n)->()"). Supported only on backends that provide a nativevectorize.
- 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 supportsignature.
- 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
vstacksemantics.- 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,BackendOpsBackendOps 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
BackendOpsmethods. Direct use ofops.torchis an explicit PyTorch-specific escape hatch.TorchOpsfollows PyTorch dtype defaults. When no dtype is provided,sanitize_dtype(None)returnstorch.get_default_dtype(). Pythoncomplexmaps totorch.complex64ortorch.complex128based 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.
- property sparse_array: Tuple[Type[Any], ...]#
Sparse array type tuple using PyTorch.
- Return type:
Tensor class accepted by this backend for sparse tensor layouts.
- 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:
- 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:
Nonefollowstorch.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:
- 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
Truebecause supported PyTorch versions provide nativevmap.
- 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)whenreturn_signis 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
logsumexpdoes not expose SciPy-styleborreturn_signparameters.
- 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:
- 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
xalong 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_stacksemantics.- 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
dstacksemantics.- 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
forloop.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
Noneif 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 returnsNone(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 reportingNoneis 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
hstacksemantics.- 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
xwithvaluesaccumulated intox[index].With
copy=Truea mutable copy is updated and returned; withcopy=Falsexis 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
xwithx[index]set tovalues.With
copy=Truea mutable copy is updated and returned; withcopy=Falsexis 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
dtypeis a complex floating type.- Parameters:
dtype (Any) – Backend or portable dtype specifier to inspect.
- Returns:
Truewhendtyperepresents 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:
dtypeitself when it is already real-valued; otherwisefloat32for complex64 andfloat64for 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.
unrollis accepted only for API parity withjax.lax.scan. WhenxsisNonean explicitlengthis required; otherwise the length is inferred from the leading axis of the first leaf ofxs.- 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
vdotsemantics.DenseLinOp.rapplyrelies 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
pyfuncelementwise, broadcasting the array arguments against one another and preserving the broadcast shape. Mirrorsnumpy.vectorize().- Parameters:
pyfunc (Callable) – Function called on scalar elements of the (broadcast) inputs.
excluded (Sequence[int] | None) – Positional argument indices passed through to
pyfuncunchanged instead of being vectorized.signature (str | None) – Generalized-ufunc signature (e.g.
"(n),(n)->()"). Supported only on backends that provide a nativevectorize.
- 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 supportsignature.
- 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
vstacksemantics.- Parameters:
arrays (Sequence[DenseArray])
- Return type:
DenseArray
- while_loop(cond_fun, body_fun, init_val)#
Run a while-loop eagerly, converting the predicate to
booleach step.- Parameters:
cond_fun (Callable[[T], bool])
body_fun (Callable[[T], T])
init_val (T)
- Return type:
T