Coverage for jaxquantum/core/sparse_bcoo.py: 94%
176 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Sparse BCOO backend for Qarray.
3Implements the JAX experimental BCOO sparse format as a Qarray storage backend.
4"""
6from __future__ import annotations
8from copy import deepcopy
10import jax.numpy as jnp
11from flax import struct
12from jax import Array
13from jax.experimental import sparse
15# QarrayImpl and friends are imported below (after qarray.py is fully loaded)
16# to match the pattern used by sparse_dia.py and avoid circular imports.
17from jaxquantum.core.qarray import DenseImpl, QarrayImpl, QarrayImplType # noqa: E402
18from jaxquantum.core.settings import SETTINGS # noqa: E402
21@struct.dataclass
22class SparseBCOOImpl(QarrayImpl):
23 """Sparse implementation using JAX experimental BCOO sparse arrays.
25 Attributes:
26 _data: The underlying ``sparse.BCOO`` array.
27 """
29 _data: sparse.BCOO
31 PROMOTION_ORDER = 1 # noqa: RUF012 — not a struct field; no annotation intentional
33 @classmethod
34 def _make(cls, data) -> "SparseBCOOImpl":
35 """Construct a ``SparseBCOOImpl`` (no sharding applied).
37 SparseBCOO has variable nnz per row and does not shard cleanly, so
38 ``_make`` is a pure pass-through. The default-sharding gate lives in
39 :meth:`from_data` instead, where users construct sharded Qarrays.
40 """
41 return cls(_data=data)
43 @classmethod
44 def from_data(cls, data) -> "SparseBCOOImpl":
45 """Wrap *data* in a new ``SparseBCOOImpl``, converting to BCOO if needed.
47 Args:
48 data: A ``sparse.BCOO`` or array-like input.
50 Returns:
51 A ``SparseBCOOImpl`` wrapping a BCOO representation of *data*.
53 Raises:
54 NotImplementedError: If a default sharding is configured. BCOO
55 does not partition cleanly across devices; convert to Dense
56 or SparseDIA, or call ``jqt.clear_default_sharding()``.
57 """
58 if SETTINGS["default_sharding"] is not None:
59 raise NotImplementedError(
60 "Sharding is not supported for SparseBCOO (variable nnz "
61 "per shard). Convert to Dense or SparseDIA, or call "
62 "`jqt.clear_default_sharding()` before constructing a "
63 "SparseBCOO Qarray."
64 )
65 return cls._make(cls._to_sparse(data))
67 def get_data(self) -> Array:
68 """Return the underlying BCOO sparse array."""
69 return self._data
71 def matmul(self, other: QarrayImpl) -> QarrayImpl:
72 """Matrix multiply ``self @ other``.
74 When *other* is a ``DenseImpl``, JAX's native BCOO @ dense path is
75 used (no self-densification). When *other* is also a
76 ``SparseBCOOImpl``, a sparse @ sparse product is performed.
78 Args:
79 other: Right-hand operand.
81 Returns:
82 A ``DenseImpl`` (sparse @ dense) or ``SparseBCOOImpl`` (sparse @
83 sparse) containing the matrix product.
84 """
85 if isinstance(other, DenseImpl):
86 return DenseImpl._make(self._data @ other._data)
87 a, b = self._coerce(other)
88 if a is not self:
89 return a.matmul(b)
90 return SparseBCOOImpl._make(self._data @ b._data)
92 def add(self, other: QarrayImpl) -> QarrayImpl:
93 """Element-wise addition ``self + other``, coercing types as needed.
95 Args:
96 other: Right-hand operand.
98 Returns:
99 A ``SparseBCOOImpl`` (both sparse) or ``DenseImpl`` (mixed) sum.
100 """
101 a, b = self._coerce(other)
102 if a is not self:
103 return a.add(b)
104 x, y = self._data, b._data
105 x, y = self._broadcast_pair(x, y)
106 if x.indices.dtype != y.indices.dtype:
107 y = sparse.BCOO((y.data, y.indices.astype(x.indices.dtype)), shape=y.shape)
108 return SparseBCOOImpl._make(x + y)
110 def sub(self, other: QarrayImpl) -> QarrayImpl:
111 """Element-wise subtraction ``self - other``, coercing types as needed.
113 Args:
114 other: Right-hand operand.
116 Returns:
117 A ``SparseBCOOImpl`` (both sparse) or ``DenseImpl`` (mixed) difference.
118 """
119 a, b = self._coerce(other)
120 if a is not self:
121 return a.sub(b)
122 x, y = self._data, b._data
123 x, y = self._broadcast_pair(x, y)
124 if x.indices.dtype != y.indices.dtype:
125 y = sparse.BCOO((y.data, y.indices.astype(x.indices.dtype)), shape=y.shape)
126 return SparseBCOOImpl._make(x - y)
128 @staticmethod
129 def _broadcast_pair(x: sparse.BCOO, y: sparse.BCOO):
130 shape = jnp.broadcast_shapes(x.shape, y.shape)
131 if x.shape != shape:
132 dims = range(len(shape) - x.ndim, len(shape))
133 x = sparse.bcoo_broadcast_in_dim(
134 x, shape=shape, broadcast_dimensions=dims
135 )
136 if y.shape != shape:
137 dims = range(len(shape) - y.ndim, len(shape))
138 y = sparse.bcoo_broadcast_in_dim(
139 y, shape=shape, broadcast_dimensions=dims
140 )
141 return x, y
143 def mul(self, scalar) -> QarrayImpl:
144 """Scalar multiplication.
146 Args:
147 scalar: Scalar value.
149 Returns:
150 A ``SparseBCOOImpl`` with each stored value multiplied by *scalar*.
151 """
152 return SparseBCOOImpl._make(scalar * self._data)
154 def dag(self) -> QarrayImpl:
155 """Conjugate transpose without densifying.
157 Transposes the last two dimensions of the BCOO array and conjugates
158 the stored values.
160 Returns:
161 A ``SparseBCOOImpl`` containing the conjugate transpose.
162 """
163 ndim = self._data.ndim
164 if ndim >= 2:
165 permutation = tuple(range(ndim - 2)) + (ndim - 1, ndim - 2)
166 transposed_data = sparse.bcoo_transpose(self._data, permutation=permutation)
167 else:
168 transposed_data = self._data
170 conjugated_data = sparse.BCOO(
171 (jnp.conj(transposed_data.data), transposed_data.indices),
172 shape=transposed_data.shape,
173 )
174 return SparseBCOOImpl._make(conjugated_data)
176 def to_dense(self) -> "DenseImpl":
177 """Convert to a ``DenseImpl`` via ``todense()``.
179 Returns:
180 A ``DenseImpl`` with the same values as this sparse array.
181 """
182 return DenseImpl._make(self._data.todense())
184 @classmethod
185 def _to_sparse(cls, data) -> sparse.BCOO:
186 """Convert *data* to a ``sparse.BCOO``, returning it unchanged if already sparse.
188 Args:
189 data: A ``sparse.BCOO`` or array-like.
191 Returns:
192 A ``sparse.BCOO`` representation of *data*.
193 """
194 if isinstance(data, sparse.BCOO):
195 return data
196 return sparse.BCOO.fromdense(data)
198 def to_sparse_bcoo(self) -> "SparseBCOOImpl":
199 """Return self (already sparse BCOO).
201 Returns:
202 This ``SparseBCOOImpl`` instance unchanged.
203 """
204 return self
206 def shape(self) -> tuple:
207 """Shape of the underlying BCOO array.
209 Returns:
210 Tuple of dimension sizes.
211 """
212 return self._data.shape
214 def dtype(self):
215 """Data type of the underlying BCOO array.
217 Returns:
218 The dtype of ``_data``.
219 """
220 return self._data.dtype
222 def frobenius_norm(self) -> float:
223 """Compute the Frobenius norm directly from stored values.
225 Returns:
226 The Frobenius norm as a scalar.
227 """
228 return jnp.sqrt(jnp.sum(jnp.abs(self._data.data) ** 2))
230 @classmethod
231 def _real(cls, data):
232 """Return a BCOO array with only the real parts of the stored values."""
233 return sparse.BCOO((jnp.real(data.data), data.indices), shape=data.shape)
235 def real(self) -> QarrayImpl:
236 """Element-wise real part.
238 Returns:
239 A ``SparseBCOOImpl`` containing the real parts of stored values.
240 """
241 return SparseBCOOImpl._make(SparseBCOOImpl._real(self._data))
243 @classmethod
244 def _imag(cls, data):
245 """Return a BCOO array with only the imaginary parts of the stored values."""
246 return sparse.BCOO((jnp.imag(data.data), data.indices), shape=data.shape)
248 def imag(self) -> QarrayImpl:
249 """Element-wise imaginary part.
251 Returns:
252 A ``SparseBCOOImpl`` containing the imaginary parts of stored values.
253 """
254 return SparseBCOOImpl._make(SparseBCOOImpl._imag(self._data))
256 @classmethod
257 def _conj(cls, data):
258 """Return a BCOO array with complex-conjugated stored values."""
259 return sparse.BCOO((jnp.conj(data.data), data.indices), shape=data.shape)
261 def conj(self) -> QarrayImpl:
262 """Element-wise complex conjugate.
264 Returns:
265 A ``SparseBCOOImpl`` containing the complex-conjugated stored values.
266 """
267 return SparseBCOOImpl._make(SparseBCOOImpl._conj(self._data))
269 @classmethod
270 def _abs(cls, data):
271 """Return a BCOO array with absolute values of stored entries."""
272 return sparse.sparsify(jnp.abs)(data)
274 def abs(self) -> QarrayImpl:
275 """Element-wise absolute value.
277 Returns:
278 A ``SparseBCOOImpl`` containing the absolute values of stored entries.
279 """
280 return SparseBCOOImpl._make(SparseBCOOImpl._abs(self._data))
282 @classmethod
283 def _eye_data(cls, n: int, dtype=None):
284 """Create an ``n x n`` identity matrix as a sparse BCOO with O(n) memory.
286 Args:
287 n: Matrix size.
288 dtype: Optional data type.
290 Returns:
291 A ``sparse.BCOO`` identity matrix of shape ``(n, n)``.
292 """
293 return sparse.eye(n, dtype=dtype)
295 @classmethod
296 def _scaled_identity(cls, n: int, scalar, dtype=None) -> SparseBCOOImpl:
297 """Create a batched scaled identity without dense broadcasting."""
298 scalar = jnp.ones((), dtype=dtype) * (jnp.asarray(scalar) + 0.0j)
299 batch_shape = scalar.shape
300 diagonal = jnp.arange(n)
301 indices = jnp.stack((diagonal, diagonal), axis=-1)
302 indices = jnp.broadcast_to(indices, (*batch_shape, n, 2))
303 values = jnp.broadcast_to(scalar[..., None], (*batch_shape, n))
304 return cls._make(
305 sparse.BCOO((values, indices), shape=(*batch_shape, n, n))
306 )
308 @classmethod
309 def can_handle_data(cls, arr) -> bool:
310 """Return True when *arr* is a ``sparse.BCOO`` array.
312 Args:
313 arr: Raw array.
315 Returns:
316 True if *arr* is a ``sparse.BCOO`` instance.
317 """
318 return isinstance(arr, sparse.BCOO)
320 @classmethod
321 def dag_data(cls, arr: sparse.BCOO) -> sparse.BCOO:
322 """Conjugate transpose for BCOO sparse arrays without densifying.
324 Args:
325 arr: A ``sparse.BCOO`` array with ``ndim >= 2``.
327 Returns:
328 A ``sparse.BCOO`` containing the conjugate transpose.
329 """
330 ndim = arr.ndim
331 permutation = tuple(range(ndim - 2)) + (ndim - 1, ndim - 2)
332 transposed = sparse.bcoo_transpose(arr, permutation=permutation)
333 return sparse.BCOO(
334 (jnp.conj(transposed.data), transposed.indices),
335 shape=transposed.shape,
336 )
338 def trace(self) -> Array:
339 """Compute the trace of the last two matrix dimensions without densifying.
341 Returns:
342 Trace value(s).
343 """
344 indices = self._data.indices
345 values = self._data.data
346 ndim = indices.shape[-1]
348 is_diag = indices[:, -2] == indices[:, -1]
350 if ndim == 2:
351 return jnp.sum(values * is_diag)
352 else:
353 batch_shape = self._data.shape[:-2]
354 B = int(jnp.prod(jnp.array(batch_shape)))
355 strides = [1]
356 for s in reversed(batch_shape[1:]):
357 strides.insert(0, strides[0] * s)
358 strides = jnp.array(strides, dtype=jnp.int32)
359 flat_batch_idx = jnp.sum(indices[:, :-2] * strides, axis=-1)
360 result = jnp.zeros(B, dtype=values.dtype).at[flat_batch_idx].add(
361 values * is_diag
362 )
363 return result.reshape(batch_shape)
365 def keep_only_diag(self) -> "SparseBCOOImpl":
366 """Zero out off-diagonal stored entries without densifying.
368 Returns:
369 A ``SparseBCOOImpl`` with only diagonal entries non-zero.
370 """
371 indices = self._data.indices
372 values = self._data.data
373 is_diag = indices[:, -2] == indices[:, -1]
374 new_values = values * is_diag
375 return SparseBCOOImpl._make(sparse.BCOO((new_values, indices), shape=self._data.shape))
377 def l2_norm_batched(self, bdims: tuple) -> Array:
378 """Compute the L2 norm per batch element without densifying.
380 Args:
381 bdims: Tuple of batch dimension sizes.
383 Returns:
384 Scalar or array of L2 norms.
385 """
386 values = self._data.data
387 indices = self._data.indices
388 n_batch_dims = len(bdims)
389 sq = jnp.abs(values) ** 2
391 if n_batch_dims == 0:
392 return jnp.sqrt(jnp.sum(sq))
393 else:
394 B = int(jnp.prod(jnp.array(bdims)))
395 strides = [1]
396 for s in reversed(bdims[1:]):
397 strides.insert(0, strides[0] * s)
398 strides = jnp.array(strides, dtype=jnp.int32)
399 flat_batch_idx = jnp.sum(indices[:, :n_batch_dims] * strides, axis=-1)
400 sum_sq = (
401 jnp.zeros(B, dtype=jnp.float64)
402 .at[flat_batch_idx]
403 .add(sq)
404 )
405 return jnp.sqrt(sum_sq).reshape(bdims)
407 def __deepcopy__(self, memo=None):
408 return SparseBCOOImpl._make(deepcopy(self._data, memo))
410 def tidy_up(self, atol) -> "SparseBCOOImpl":
411 """Zero out stored values whose real or imaginary magnitude is below *atol*.
413 Args:
414 atol: Absolute tolerance threshold.
416 Returns:
417 A new ``SparseBCOOImpl`` with small values zeroed.
418 """
419 values = self._data.data
420 re = jnp.real(values)
421 im = jnp.imag(values)
422 new_values = re * (jnp.abs(re) > atol) + 1j * im * (jnp.abs(im) > atol)
423 return SparseBCOOImpl._make(
424 sparse.BCOO((new_values, self._data.indices), shape=self._data.shape)
425 )
427 def kron(self, other: "QarrayImpl") -> "QarrayImpl":
428 """Kronecker product using ``sparsify(jnp.kron)`` — stays sparse.
430 Args:
431 other: Right-hand operand.
433 Returns:
434 A ``SparseBCOOImpl`` containing the Kronecker product when both
435 operands are sparse; a ``DenseImpl`` when types differ.
436 """
437 a, b = self._coerce(other)
438 if a is not self:
439 return a.kron(b)
440 sparse_kron = sparse.sparsify(jnp.kron)
441 return SparseBCOOImpl._make(sparse_kron(self._data, b._data))
444# Register with the enum registry
445QarrayImplType.register(SparseBCOOImpl, QarrayImplType.SPARSE_BCOO)
447__all__ = ["SparseBCOOImpl"]