from __future__ import annotations
from typing import TYPE_CHECKING, Any, Sequence, Literal, Tuple, Callable, Optional, Type, cast
from warnings import warn
from .._family import BackendFamily
from .._ops import BackendOps
from ..numpy import NumpyOps
from ...types import DenseArray, ArrayLike, SparseArray, DType, Index, X, T, Y, R, Carry
try: # private; see JaxOps.is_tracing
from jax._src import core as _jax_core
except ImportError:
_jax_core = None
if TYPE_CHECKING:
from jax import Array
NativeArray = Array
else:
NativeArray = DenseArray
[docs]
class JaxOps(BackendOps[NativeArray]):
"""
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.
"""
import jax as _jax
import jax.numpy as _jnp
import jax.experimental.sparse as _jsparse
# Concrete library handles exposed as ``Any`` so the portable protocols can
# flow into typed JAX calls without per-boundary casts; mirrors the base
# ``xp: ClassVar[Any]`` design.
jax: Any = _jax
jnp: Any = _jnp
jsparse: Any = _jsparse
xp = jnp
_family = BackendFamily.jax.value
_allow_sparse = True
[docs]
@classmethod
def install_pytree_protocol(cls) -> None:
"""Register SpaceCore containers as JAX pytree nodes.
JAX's ``register_pytree_node_class`` consumes exactly the protocol
:class:`~spacecore.backend._container.PyTreeNode` defines — ``tree_flatten``
returning ``(children, aux)`` and a ``tree_unflatten(aux, children)``
classmethod — so it *is* the registrar and no translation is needed.
"""
import jax
from .._container import registry
registry.register_backend("jax", jax.tree_util.register_pytree_node_class)
[docs]
def sanitize_dtype(self, dtype: DType | None) -> DType:
"""Normalize a dtype for JAX, rejecting dtypes JAX would silently canonicalize under the active x64 setting."""
# Default: float64 if x64 is enabled, else float32 with a warning.
x64_enabled = bool(self.jax.config.read("jax_enable_x64"))
if dtype is None:
if not x64_enabled:
warn(
"jax_enable_x64 is set to False, so default JAX dtype is set to float32. "
"If you need float64, run `jax.config.update('jax_enable_x64', True)`.",
UserWarning,
)
return self.jnp.float32
return self.jnp.float64
try:
dt = self.jnp.dtype(dtype)
except Exception as e:
raise TypeError(f"Invalid dtype specifier for JAX: {dtype!r}.") from e
# Ensure dtype is actually usable on this backend/device
try:
self.jnp.empty((), dtype=dt)
except Exception as e:
raise TypeError(
f"Dtype {dt!r} is not supported by the active JAX backend/device."
) from e
# Forbid implicit coercion under current JAX configuration
dt_canon = self.jax.dtypes.canonicalize_dtype(dt)
if dt_canon != dt:
raise TypeError(
f"Dtype {dt} is not permitted under current JAX configuration: "
f"it would be canonicalized to {dt_canon}. "
f"(jax_enable_x64={x64_enabled!r})"
)
return dt
@property
def dense_array(self) -> Type[Any]:
"""
Dense array type using JAX.
Returns
-------
Concrete dense array class accepted by this backend.
See:
https://docs.jax.dev/en/latest/jax.Array.html
"""
return self.jax.Array
@property
def sparse_array(self) -> Tuple[Type[Any], ...]:
"""
Sparse array type tuple using JAX.
Returns
-------
Concrete sparse array classes accepted by this backend, or None.
See:
https://docs.jax.dev/en/latest/jax.experimental.sparse.html
"""
return (self.jsparse.BCOO, self.jsparse.BCSR)
[docs]
def assparse(
self,
x: Any,
*,
format: Literal["bcoo", "bcsr"] = "bcoo",
index_dtype: DType | None = None,
nse: int | None = None,
dtype: DType | None = None,
) -> SparseArray:
"""Convert to a JAX sparse array (BCOO/BCSR); SciPy inputs go through ``from_scipy_sparse``."""
import scipy.sparse as sps
self._reject_complex_to_real(x, dtype, operation="assparse")
# Already a JAX sparse array: cast only.
if self.is_sparse(x):
if dtype is not None and self.get_dtype(x) != self.sanitize_dtype(dtype):
return x.astype(self.sanitize_dtype(dtype))
return x
classes = {"bcoo": self.jsparse.BCOO, "bcsr": self.jsparse.BCSR}
if format not in classes:
raise ValueError(f"Unknown sparse format: {format!r}")
kwargs = {k: v for k, v in (("index_dtype", index_dtype), ("nse", nse)) if v is not None}
# SciPy sparse input converts directly, without densifying; dense input goes through fromdense.
if sps.issparse(x):
out = classes[format].from_scipy_sparse(x, **kwargs)
return out.astype(self.sanitize_dtype(dtype)) if dtype is not None else out
return classes[format].fromdense(self.asarray(x, dtype=dtype), **kwargs)
[docs]
def sparse_matmul(self, a: SparseArray, b: DenseArray) -> NativeArray:
"""Multiply a JAX sparse array by a dense one, returning a dense ``jax.Array``."""
return a @ b
[docs]
def vmap(
self,
fn: Callable,
in_axes: int | Sequence[int | None] | None = 0,
out_axes: int | Sequence[int | None] | None = 0,
) -> Callable:
"""Vectorize a function using ``jax.vmap``."""
return self.jax.vmap(fn, in_axes=in_axes, out_axes=out_axes)
[docs]
def is_tracing(self) -> bool:
"""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.
"""
clean = getattr(_jax_core, "trace_state_clean", None)
return clean is not None and not clean()
@property
def has_native_vmap(self) -> bool:
"""Return ``True`` because JAX provides native ``vmap``."""
return True
[docs]
def logsumexp(
self,
a: DenseArray,
axis: int | Sequence[int] | None = None,
b: DenseArray | None = None,
keepdims: bool = False,
return_sign: bool = False,
where: DenseArray | None = None,
) -> NativeArray | Tuple[NativeArray, NativeArray]:
"""
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
"""
return self.jax.scipy.special.logsumexp(
a, axis=axis, b=b, keepdims=keepdims, return_sign=return_sign, where=where
)
[docs]
def index_set(
self, x: DenseArray, index: Index, values: ArrayLike | float | complex,
*, copy: bool = True,
) -> NativeArray:
"""Return ``x`` with ``x[index] = values`` via ``x.at[index].set``; ``copy=False`` raises, as JAX arrays are immutable."""
if not copy:
raise NotImplementedError("JAX arrays are immutable; copy=False is not supported.")
return cast(Any, x).at[index].set(values)
[docs]
def ix_(self, *args: Any) -> Any:
r"""
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 <https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ix_.html>`_
"""
return self.jnp.ix_(*args)
[docs]
def fori_loop(
self,
lower: int,
upper: int,
body_fun: Callable[[int, T], T],
init_val: T,
*,
unroll: int | bool | None = None,
) -> T:
"""Run ``jax.lax.fori_loop``."""
return self.jax.lax.fori_loop(lower, upper, body_fun, init_val, unroll=unroll)
[docs]
def while_loop(
self,
cond_fun: Callable[[T], bool],
body_fun: Callable[[T], T],
init_val: T,
) -> T:
"""Run ``jax.lax.while_loop``."""
return self.jax.lax.while_loop(cond_fun, body_fun, init_val)
[docs]
def scan(
self,
f: Callable[[Carry, X], Tuple[Carry, Y]],
init: Carry,
xs: X,
length: Optional[int] = None,
reverse: bool = False,
unroll: int = 1,
_split_transpose: bool = False,
) -> Tuple[Carry, Y]:
"""Run ``jax.lax.scan``; carries and outputs may be pytrees."""
return self.jax.lax.scan(
f,
init,
xs,
length=length,
reverse=reverse,
unroll=unroll,
_split_transpose=_split_transpose,
)
[docs]
def cond(
self,
pred: bool,
true_fun: Callable[[T], R],
false_fun: Callable[[T], R],
*operands: Any,
) -> R:
"""Run ``jax.lax.cond``: both branches are traced, not chosen eagerly in Python."""
return self.jax.lax.cond(pred, true_fun, false_fun, *operands)
[docs]
def index_add(
self, x: DenseArray, index: Index, values: DenseArray, *, copy: bool = True
) -> NativeArray:
"""Return ``x`` with ``values`` added at ``index`` via ``x.at[index].add``; repeated indices accumulate, and ``copy=False`` raises."""
if not copy:
raise NotImplementedError("JAX arrays are immutable; copy=False is not supported.")
return cast(Any, x).at[index].add(values)
[docs]
def allclose_sparse(
self,
a: SparseArray,
b: SparseArray,
rtol: float = 1e-5,
atol: float = 1e-8,
) -> bool:
"""Compare two JAX sparse arrays within tolerances, via SciPy sparse."""
self._require_two_sparse(a, b)
# Compare on the host with the SciPy implementation.
np_ops = NumpyOps()
a_sp = self._to_scipy_sparse(np_ops, a)
b_sp = self._to_scipy_sparse(np_ops, b)
return np_ops.allclose_sparse(a_sp, b_sp, rtol=rtol, atol=atol)
def _to_scipy_sparse(self, np_ops: NumpyOps, x: SparseArray):
# Normalize to 2-D unbatched BCOO, then build a SciPy COO from its indices.
bcoo: Any = x
if isinstance(x, self.jsparse.BCSR):
bcoo = bcoo.to_bcoo()
if isinstance(bcoo, self.jsparse.BCOO):
bcoo = bcoo.sum_duplicates(remove_zeros=False)
if bcoo.n_batch != 0 or bcoo.n_dense != 0 or bcoo.n_sparse != 2:
raise NotImplementedError(
"_to_scipy_sparse supports only 2D unbatched sparse arrays."
)
row = bcoo.indices[:, 0]
col = bcoo.indices[:, 1]
data = bcoo.data
return np_ops.sp.sparse.coo_array((data, (row, col)), shape=bcoo.shape)
raise TypeError(f"Unsupported sparse type: {type(x)!r}")
