Coverage for jaxquantum/core/sparse_dia.py: 93%
263 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Sparse diagonal (SparseDIA) backend for Qarray.
3Stores only the diagonal *values* of a matrix, making quantum operators with
4small numbers of non-zero diagonals (annihilation, creation, number, Kerr…)
5far cheaper than Dense or BCOO:
7 * Memory: O(d * n) where d = number of stored diagonals, n = matrix size
8 * No index arrays (unlike BCOO which stores (row, col) per non-zero)
9 * ``_offsets`` is *static* Python metadata (pytree_node=False), so JAX
10 unrolls all loops over diagonals at compile time — only static slices,
11 no dynamic indexing or scatter/gather.
13Padding convention (Convention A):
15 For diagonal at offset k (k ≥ 0):
16 diags[..., i, j] = A[j-k, j] for j ∈ [k, n-1]; zeros at [0:k]
17 For diagonal at offset k (k < 0):
18 diags[..., i, j] = A[j-k, j] for j ∈ [0, n+k-1]; zeros at [n+k:]
20Unified access formula (holds for any k, out-of-range slots are zero):
21 A[i, i+k] = diags[..., diag_idx, i+k]
23This makes every matrix operation a set of aligned slice multiplications.
25Some improvements (_dia_slice helper, integer matrix power, diagonal-range pruning,
26offset detection) were identified by studying the dynamiqs library
27(https://github.com/dynamiqs/dynamiqs).
28"""
30from __future__ import annotations
32import numpy as np
33from copy import deepcopy
34from typing import TYPE_CHECKING
36import jax.numpy as jnp
37from flax import struct
38from jax import Array
40if TYPE_CHECKING:
41 from jaxquantum.core.qarray import DenseImpl, QarrayImplType
42 from jaxquantum.core.sparse_bcoo import SparseBCOOImpl
44from jaxquantum.core.settings import _maybe_shard
47# ---------------------------------------------------------------------------
48# Slice helper
49# ---------------------------------------------------------------------------
51def _dia_slice(k: int) -> slice:
52 """Slice selecting the valid data positions for diagonal offset k.
54 For k ≥ 0: valid data lives at column indices [k, n), so slice(k, None).
55 For k < 0: valid data lives at column indices [0, n+k), so slice(None, k).
57 The complementary slice (for the result rows / left-operand columns) is
58 always ``_dia_slice(-k)``.
59 """
60 return slice(k, None) if k >= 0 else slice(None, k)
63# ---------------------------------------------------------------------------
64# Raw data container
65# ---------------------------------------------------------------------------
67@struct.dataclass
68class SparseDiaData:
69 """Lightweight pytree-compatible container for sparse-diagonal raw data.
71 Returned by ``SparseDiaImpl.get_data()`` and consumed by
72 ``SparseDiaImpl.from_data()``. Registered as a JAX pytree via Flax's
73 ``@struct.dataclass``; ``offsets`` is *not* a pytree leaf (it is static
74 compile-time metadata).
76 Attributes:
77 offsets: Static tuple of diagonal offsets (pytree_node=False).
78 diags: JAX array of shape (*batch, n_diags, n) containing the
79 padded diagonal values.
80 """
82 offsets: tuple = struct.field(pytree_node=False)
83 diags: Array
85 # Class-level marker (not a dataclass field — no type annotation).
86 # Used by DenseImpl.can_handle_data to exclude SparseDiaData without
87 # a direct import (which would be circular).
88 _is_sparse_dia = True
90 @property
91 def shape(self) -> tuple:
92 """Shape of the represented square matrix (*batch, n, n)."""
93 n = self.diags.shape[-1]
94 return (*self.diags.shape[:-2], n, n)
96 @property
97 def dtype(self):
98 """Dtype of the stored diagonal values."""
99 return self.diags.dtype
101 def __mul__(self, scalar):
102 return SparseDiaData(offsets=self.offsets, diags=self.diags * scalar)
104 def __rmul__(self, scalar):
105 return SparseDiaData(offsets=self.offsets, diags=scalar * self.diags)
107 def __getitem__(self, index):
108 """Index into the batch dimension(s), preserving offsets."""
109 return SparseDiaData(offsets=self.offsets, diags=self.diags[index])
111 def __len__(self):
112 """Number of elements along the leading batch dimension."""
113 return self.shape[0]
115 def reshape(self, *new_shape):
116 """Reshape batch dimensions while preserving diagonal structure.
118 ``new_shape`` must end with ``(N, N)`` (the matrix dims are unchanged).
119 Only the leading batch dims are reshaped.
120 """
121 if len(new_shape) == 1 and isinstance(new_shape[0], (tuple, list)):
122 new_shape = new_shape[0]
123 new_batch = new_shape[:-2]
124 n = self.diags.shape[-1]
125 new_diags = self.diags.reshape(*new_batch, len(self.offsets), n)
126 return SparseDiaData(offsets=self.offsets, diags=new_diags)
128 def __matmul__(self, other):
129 """SparseDIA @ dense → dense (used by mesolve ODE RHS)."""
130 # _sparsedia_matmul_dense is defined later in this module; Python
131 # resolves the name at call time so forward reference is fine.
132 return _sparsedia_matmul_dense(self.offsets, self.diags, other)
134 def __rmatmul__(self, other):
135 """dense @ SparseDIA → dense (used by mesolve ODE RHS)."""
136 return _sparsedia_rmatmul_dense(other, self.offsets, self.diags)
139# ---------------------------------------------------------------------------
140# Helper: dense → SparseDIA conversion
141# ---------------------------------------------------------------------------
143def _dense_to_sparsedia(arr: np.ndarray) -> tuple[tuple, np.ndarray]:
144 """Extract non-zero diagonal offsets and padded values from a dense array.
146 Uses the first batch element (if batched) to detect which diagonals are
147 non-zero. Safe to call outside JIT because *arr* must be a concrete
148 numpy / JAX array.
150 Args:
151 arr: Dense array of shape (*batch, n, n).
153 Returns:
154 Tuple of (offsets, diags) where:
155 - offsets is a sorted tuple of integer offsets
156 - diags is a numpy array of shape (*batch, n_diags, n)
157 """
158 n = arr.shape[-1]
159 batch_shape = arr.shape[:-2]
161 # Union non-zero diagonals across all batch elements via a single mask + nonzero call.
162 arr_np = np.asarray(arr)
163 flat_np = arr_np.reshape(-1, n, n)
164 union_mask = np.any(flat_np != 0, axis=0) # (n, n): True where any batch elem is non-zero
165 r, c = np.nonzero(union_mask)
166 offsets = tuple(sorted(set((c - r).tolist()))) if len(r) > 0 else (0,)
168 diags = np.zeros((*batch_shape, len(offsets), n), dtype=arr_np.dtype)
169 for i, k in enumerate(offsets):
170 # np.diagonal returns shape (*batch, n-|k|)
171 d = np.diagonal(arr_np, offset=k, axis1=-2, axis2=-1)
172 lo = max(k, 0)
173 hi = n - max(-k, 0)
174 diags[..., i, lo:hi] = d
176 return offsets, diags
179# ---------------------------------------------------------------------------
180# SparseDiaImpl
181# ---------------------------------------------------------------------------
183from jaxquantum.core.qarray import QarrayImpl, DenseImpl, QarrayImplType # noqa: E402
186@struct.dataclass
187class SparseDiaImpl(QarrayImpl):
188 """Sparse-diagonal backend storing only diagonal values.
190 Data layout::
192 _offsets : tuple[int, ...] — static (pytree_node=False)
193 _diags : Array[*batch, n_diags, n] — JAX-traced values
195 For offset k, the convention is:
196 * k ≥ 0 : valid data at ``_diags[..., i, k:]``, zeros at ``[0:k]``
197 * k < 0 : valid data at ``_diags[..., i, :n+k]``, zeros at ``[n+k:]``
199 In both cases: ``A[row, row+k] = _diags[..., i, row+k]``
200 """
202 _offsets: tuple = struct.field(pytree_node=False)
203 _diags: Array
205 PROMOTION_ORDER = 0 # noqa: RUF012 — not a struct field
207 # ------------------------------------------------------------------
208 # Construction
209 # ------------------------------------------------------------------
211 # NOTE (potential future optimization):
212 # `_diags` has shape (*batch, n_diags, n) where n_diags is typically 1-5.
213 # Sharding the n axis aligns reasonably with dense ops, but for tightly
214 # coupled SparseDIA-only chains it may be cheaper to *always replicate*
215 # `_diags` (small relative to dense (n, n) storage) and rely on the dense
216 # state's sharding to drive parallelism. Benchmark before committing — this
217 # is a behavior change for users who want SparseDIA sharding intentionally.
218 @classmethod
219 def _make(cls, offsets: tuple, diags: Array) -> "SparseDiaImpl":
220 """Construct a ``SparseDiaImpl``, applying the configured default sharding.
222 All internal construction sites route through this so every Qarray
223 (including intermediates produced by ``add``, ``matmul``, ``kron``,
224 etc.) satisfies the user's sharding invariant.
225 """
226 return cls(_offsets=offsets, _diags=_maybe_shard(diags))
228 @classmethod
229 def from_data(cls, data) -> "SparseDiaImpl":
230 """Wrap *data* in a new ``SparseDiaImpl``.
232 Accepts either a :class:`SparseDiaData` container (direct wrap) or
233 a dense array-like (auto-detect non-zero diagonals via numpy, safe
234 to call before JIT).
236 Args:
237 data: A :class:`SparseDiaData` or dense array of shape
238 (*batch, n, n).
240 Returns:
241 A new ``SparseDiaImpl`` instance.
242 """
243 if isinstance(data, SparseDiaData):
244 return cls._make(data.offsets, data.diags)
245 offsets, diags_np = _dense_to_sparsedia(np.asarray(data))
246 return cls._make(offsets, jnp.array(diags_np))
248 @classmethod
249 def from_diags(cls, offsets: tuple, diags: Array) -> "SparseDiaImpl":
250 """Directly construct from sorted offsets and padded diagonal array.
252 This is the preferred factory when diagonal structure is known in
253 advance (e.g., when building ``destroy`` or ``create`` operators).
255 Args:
256 offsets: Tuple of integer diagonal offsets (need not be sorted;
257 will be sorted internally).
258 diags: JAX array of shape (*batch, n_diags, n) with padded
259 diagonal values matching *offsets*.
261 Returns:
262 A new ``SparseDiaImpl`` instance.
263 """
264 return cls._make(tuple(sorted(offsets)), diags)
266 # ------------------------------------------------------------------
267 # QarrayImpl abstract methods
268 # ------------------------------------------------------------------
270 def get_data(self) -> SparseDiaData:
271 """Return a :class:`SparseDiaData` container with the raw diagonal data."""
272 return SparseDiaData(offsets=self._offsets, diags=self._diags)
274 def shape(self) -> tuple:
275 """Shape of the represented square matrix (including batch dims)."""
276 n = self._diags.shape[-1]
277 return (*self._diags.shape[:-2], n, n)
279 def dtype(self):
280 """Dtype of the stored diagonal values."""
281 return self._diags.dtype
283 def __deepcopy__(self, memo=None):
284 return SparseDiaImpl._make(deepcopy(self._offsets), self._diags)
286 # ------------------------------------------------------------------
287 # Arithmetic
288 # ------------------------------------------------------------------
290 def mul(self, scalar) -> "SparseDiaImpl":
291 """Scalar multiplication — scales all diagonal values."""
292 return SparseDiaImpl._make(self._offsets, scalar * self._diags)
294 def neg(self) -> "SparseDiaImpl":
295 """Negation."""
296 return SparseDiaImpl._make(self._offsets, -self._diags)
298 def add(self, other: QarrayImpl) -> QarrayImpl:
299 """Element-wise addition.
301 SparseDIA + SparseDIA stays SparseDIA (union of offsets, static).
302 Otherwise coerces to the higher-order type.
303 """
304 if isinstance(other, SparseDiaImpl):
305 return _sparsedia_add(self, other)
306 a, b = self._coerce(other)
307 if a is not self:
308 return a.add(b)
309 return a.add(b)
311 def sub(self, other: QarrayImpl) -> QarrayImpl:
312 """Element-wise subtraction."""
313 if isinstance(other, SparseDiaImpl):
314 return _sparsedia_add(self, other, subtract=True)
315 a, b = self._coerce(other)
316 if a is not self:
317 return a.sub(b)
318 return a.sub(b)
320 def matmul(self, other: QarrayImpl) -> QarrayImpl:
321 """Matrix multiplication.
323 * SparseDIA @ SparseDIA → SparseDIA (O(d₁·d₂·n))
324 * SparseDIA @ Dense → Dense (O(d·n²), no densification of self)
325 * Others → coerce then delegate
326 """
327 if isinstance(other, DenseImpl):
328 return DenseImpl._make(_sparsedia_matmul_dense(
329 self._offsets, self._diags, other._data
330 ))
331 if isinstance(other, SparseDiaImpl):
332 offsets, diags = _sparsedia_matmul_sparsedia(
333 self._offsets, self._diags,
334 other._offsets, other._diags,
335 )
336 return SparseDiaImpl._make(offsets, diags)
337 a, b = self._coerce(other)
338 if a is not self:
339 return a.matmul(b)
340 return a.matmul(b)
342 def dag(self) -> "SparseDiaImpl":
343 """Conjugate transpose without densification.
345 Negates every offset and rearranges the stored values so that the
346 padding convention remains consistent.
347 """
348 new_offsets = tuple(-k for k in self._offsets)
349 new_diags = jnp.zeros_like(self._diags)
350 for i, k in enumerate(self._offsets):
351 s = _dia_slice(k) # valid data slice for offset k
352 sm = _dia_slice(-k) # valid data slice for offset -k (the new position)
353 new_diags = new_diags.at[..., i, sm].set(jnp.conj(self._diags[..., i, s]))
354 return SparseDiaImpl._make(new_offsets, new_diags)
356 def kron(self, other: QarrayImpl) -> QarrayImpl:
357 """Kronecker product.
359 SparseDIA ⊗ SparseDIA stays SparseDIA: output offset for pair
360 (kA, kB) is ``kA * m + kB`` where m = dim(B). Fully vectorised —
361 no loops at JAX level.
362 """
363 if isinstance(other, SparseDiaImpl):
364 return _sparsedia_kron(self, other)
365 a, b = self._coerce(other)
366 if a is not self:
367 return a.kron(b)
368 return a.kron(b)
370 def tidy_up(self, atol) -> "SparseDiaImpl":
371 """Zero diagonal values whose magnitude is below *atol*."""
372 diags = self._diags
373 real_part = jnp.where(jnp.abs(jnp.real(diags)) < atol, 0.0, jnp.real(diags))
374 if jnp.issubdtype(diags.dtype, jnp.complexfloating):
375 imag_part = jnp.where(jnp.abs(jnp.imag(diags)) < atol, 0.0, jnp.imag(diags))
376 new_diags = (real_part + 1j * imag_part).astype(diags.dtype)
377 else:
378 new_diags = real_part.astype(diags.dtype)
379 return SparseDiaImpl._make(self._offsets, new_diags)
381 # ------------------------------------------------------------------
382 # Conversions
383 # ------------------------------------------------------------------
385 def to_dense(self) -> "DenseImpl":
386 """Convert to a ``DenseImpl`` by summing diagonal contributions."""
387 n = self._diags.shape[-1]
388 batch_shape = self._diags.shape[:-2]
389 result = jnp.zeros((*batch_shape, n, n), dtype=self._diags.dtype)
390 for i, k in enumerate(self._offsets):
391 s = _dia_slice(k)
392 length = n - abs(k)
393 if length <= 0:
394 continue
395 vals = self._diags[..., i, s]
396 row_idx = jnp.arange(length) + max(-k, 0)
397 col_idx = row_idx + k
398 result = result.at[..., row_idx, col_idx].set(vals)
399 return DenseImpl._make(result)
401 def to_sparse_bcoo(self) -> "SparseBCOOImpl":
402 """Convert to a ``SparseBCOOImpl`` (BCOO) via dense."""
403 return self.to_dense().to_sparse_bcoo()
405 def to_sparse_dia(self) -> "SparseDiaImpl":
406 """Return self (already SparseDIA)."""
407 return self
409 # ------------------------------------------------------------------
410 # Class-method interface
411 # ------------------------------------------------------------------
413 @classmethod
414 def _eye_data(cls, n: int, dtype=None):
415 """Return an n×n identity as a dense JAX array."""
416 return jnp.eye(n, dtype=dtype)
418 @classmethod
419 def _scaled_identity(cls, n: int, scalar, dtype=None) -> "SparseDiaImpl":
420 """Create a batched scaled identity without dense storage."""
421 scalar = jnp.ones((), dtype=dtype) * (jnp.asarray(scalar) + 0.0j)
422 diags = jnp.broadcast_to(
423 scalar[..., None, None],
424 (*scalar.shape, 1, n),
425 )
426 return cls._make((0,), diags)
428 @classmethod
429 def can_handle_data(cls, arr) -> bool:
430 """Return True only for :class:`SparseDiaData` objects."""
431 return isinstance(arr, SparseDiaData)
433 @classmethod
434 def dag_data(cls, arr: SparseDiaData) -> SparseDiaData:
435 """Conjugate transpose of raw :class:`SparseDiaData` without densification."""
436 impl = SparseDiaImpl._make(arr.offsets, arr.diags)
437 result = impl.dag()
438 return result.get_data()
440 # ------------------------------------------------------------------
441 # Extra sparse-native methods (no densification)
442 # ------------------------------------------------------------------
444 def trace(self):
445 """Compute trace directly from the main diagonal (offset 0).
447 Returns:
448 Scalar trace (sum of main diagonal values).
449 """
450 if 0 in self._offsets:
451 i = self._offsets.index(0)
452 return jnp.sum(self._diags[..., i, :], axis=-1)
453 return jnp.zeros(self._diags.shape[:-2], dtype=self._diags.dtype)
455 def frobenius_norm(self):
456 """Frobenius norm computed directly from stored diagonal values."""
457 return jnp.sqrt(jnp.sum(jnp.abs(self._diags) ** 2))
459 def real(self) -> "SparseDiaImpl":
460 """Element-wise real part of stored values."""
461 return SparseDiaImpl._make(
462 self._offsets,
463 jnp.real(self._diags).astype(self._diags.dtype),
464 )
466 def imag(self) -> "SparseDiaImpl":
467 """Element-wise imaginary part of stored values."""
468 return SparseDiaImpl._make(
469 self._offsets,
470 jnp.imag(self._diags).astype(self._diags.dtype),
471 )
473 def conj(self) -> "SparseDiaImpl":
474 """Element-wise complex conjugate of stored values."""
475 return SparseDiaImpl._make(self._offsets, jnp.conj(self._diags))
477 def powm(self, n: int) -> "SparseDiaImpl":
478 """Integer matrix power staying SparseDIA via binary exponentiation.
480 Uses O(log n) SparseDIA @ SparseDIA multiplications rather than
481 densifying. A^0 returns the identity operator.
483 Args:
484 n: Non-negative integer exponent.
486 Returns:
487 A ``SparseDiaImpl`` equal to this matrix raised to the *n*-th power.
489 Raises:
490 ValueError: If *n* is negative.
491 """
492 if n < 0:
493 raise ValueError("powm requires n >= 0")
494 if n == 0:
495 size = self._diags.shape[-1]
496 eye_diags = jnp.ones((*self._diags.shape[:-2], 1, size), dtype=self._diags.dtype)
497 return SparseDiaImpl._make((0,), eye_diags)
498 if n == 1:
499 return self
500 half = self.powm(n // 2)
501 squared = half.matmul(half) # SparseDIA @ SparseDIA → SparseDIA
502 return squared if n % 2 == 0 else self.matmul(squared)
505# ---------------------------------------------------------------------------
506# Pure-function helpers (operate on raw arrays, no QarrayImpl wrapping)
507# ---------------------------------------------------------------------------
509def _sparsedia_add(
510 a: SparseDiaImpl,
511 b: SparseDiaImpl,
512 subtract: bool = False,
513) -> SparseDiaImpl:
514 """Add (or subtract) two SparseDiaImpl matrices, preserving SparseDIA format.
516 Computes the union of offsets at Python level (static), then copies /
517 sums the corresponding diagonal arrays.
518 """
519 n = a._diags.shape[-1]
520 batch_shape = jnp.broadcast_shapes(a._diags.shape[:-2], b._diags.shape[:-2])
522 out_offsets = tuple(sorted(set(a._offsets) | set(b._offsets)))
523 out_dtype = jnp.result_type(a._diags.dtype, b._diags.dtype)
524 out_diags = jnp.zeros((*batch_shape, len(out_offsets), n), dtype=out_dtype)
526 a_idx = {k: i for i, k in enumerate(a._offsets)}
527 b_idx = {k: i for i, k in enumerate(b._offsets)}
529 for oi, k in enumerate(out_offsets):
530 val = jnp.zeros((*batch_shape, n), dtype=out_dtype)
531 if k in a_idx:
532 val = val + a._diags[..., a_idx[k], :]
533 if k in b_idx:
534 sign = -1 if subtract else 1
535 val = val + sign * b._diags[..., b_idx[k], :]
536 out_diags = out_diags.at[..., oi, :].set(val)
538 return SparseDiaImpl._make(out_offsets, out_diags)
541def _sparsedia_matmul_dense(
542 offsets: tuple,
543 diags: Array,
544 B: Array,
545) -> Array:
546 """Compute (SparseDIA) @ (dense matrix) → dense, without densifying the LHS.
548 For each stored diagonal at offset k:
549 result[..., row_range, :] += diags[..., i, valid_slice, None] * B[..., valid_slice, :]
551 Complexity: O(d * n * m) where n×n is the operator and n×m is B.
553 Args:
554 offsets: Static tuple of diagonal offsets for the LHS.
555 diags: JAX array of shape (*batch, n_diags, n).
556 B: Dense right-hand side of shape (*batch, n, m).
558 Returns:
559 Dense product of shape (*batch, n, m).
560 """
561 # n = diags.shape[-1]
562 batch_shape = jnp.broadcast_shapes(diags.shape[:-2], B.shape[:-2])
563 result = jnp.zeros(
564 (*batch_shape, B.shape[-2], B.shape[-1]),
565 dtype=jnp.result_type(diags.dtype, B.dtype),
566 )
567 for i, k in enumerate(offsets):
568 s = _dia_slice(k) # valid column slice for diagonal k
569 sm = _dia_slice(-k) # corresponding row slice for the result
570 result = result.at[..., sm, :].add(diags[..., i, s, None] * B[..., s, :])
571 return result
574def _sparsedia_rmatmul_dense(
575 B: Array,
576 offsets: tuple,
577 diags: Array,
578) -> Array:
579 """Compute (dense matrix) @ (SparseDIA) → dense.
581 For each stored diagonal at offset k:
582 C[..., :, k:] += B[..., :, :n-k] * diags[..., i, k:][..., None, :] (k ≥ 0)
583 C[..., :, :n-m] += B[..., :, m:] * diags[..., i, :n-m][..., None, :] (k < 0)
585 Complexity: O(d * n * p) where n×n is the operator and p×n is B.
587 Args:
588 B: Dense left-hand side of shape (*batch, p, n).
589 offsets: Static tuple of diagonal offsets for the RHS.
590 diags: JAX array of shape (*batch, n_diags, n).
592 Returns:
593 Dense product of shape (*batch, p, n).
594 """
595 # n = diags.shape[-1]
596 batch_shape = jnp.broadcast_shapes(diags.shape[:-2], B.shape[:-2])
597 result = jnp.zeros(
598 (*batch_shape, B.shape[-2], B.shape[-1]),
599 dtype=jnp.result_type(diags.dtype, B.dtype),
600 )
601 for i, k in enumerate(offsets):
602 s = _dia_slice(k) # valid column slice for diagonal k
603 sm = _dia_slice(-k) # complementary slice for B columns / result columns
604 result = result.at[..., :, s].add(B[..., :, sm] * diags[..., i, s][..., None, :])
605 return result
608def _sparsedia_matmul_sparsedia(
609 left_offsets: tuple,
610 left_diags: Array,
611 right_offsets: tuple,
612 right_diags: Array,
613) -> tuple[tuple, Array]:
614 """Compute (SparseDIA) @ (SparseDIA) → SparseDIA.
616 Derivation uses the unified access formula A[i, i+k] = diags[i+k].
617 For the contribution of diagonal pair (k1, k2) to output at kout = k1+k2:
619 out_diag[j + k2] += left_diag[j] * right_diag[j + k2]
621 This is an aligned slice multiply, handled by the sign of k2:
622 k2 ≥ 0: out.at[kout:].add( left[:n-k2] * right[k2:] )
623 k2 < 0: out.at[:n+k2].add( left[-k2:] * right[:n+k2] )
625 Complexity: O(d1 * d2 * n).
627 Args:
628 left_offsets: Static offsets for the LHS matrix.
629 left_diags: JAX array (*batch, d1, n).
630 right_offsets: Static offsets for the RHS matrix.
631 right_diags: JAX array (*batch, d2, n).
633 Returns:
634 Tuple of (out_offsets, out_diags).
635 """
636 n = left_diags.shape[-1]
637 batch_shape = jnp.broadcast_shapes(
638 left_diags.shape[:-2], right_diags.shape[:-2]
639 )
641 # Pre-filter output offsets: diagonal pairs where |k1+k2| >= n are zero.
642 out_offset_set = sorted(
643 {k1 + k2 for k1 in left_offsets for k2 in right_offsets if abs(k1 + k2) < n}
644 )
645 if not out_offset_set:
646 out_offset_set = [0]
647 out_offset_idx = {k: i for i, k in enumerate(out_offset_set)}
648 out_diags = jnp.zeros(
649 (*batch_shape, len(out_offset_set), n),
650 dtype=jnp.result_type(left_diags.dtype, right_diags.dtype),
651 )
653 for li, k1 in enumerate(left_offsets):
654 for ri, k2 in enumerate(right_offsets):
655 kout = k1 + k2
656 if abs(kout) >= n:
657 continue
658 oi = out_offset_idx[kout]
659 s = _dia_slice(k2) # valid column slice for right diagonal k2
660 sm = _dia_slice(-k2) # complementary slice for left diagonal
661 contribution = left_diags[..., li, sm] * right_diags[..., ri, s]
662 out_diags = out_diags.at[..., oi, s].add(contribution)
664 return tuple(out_offset_set), out_diags
667def _sparsedia_kron(a: SparseDiaImpl, b: SparseDiaImpl) -> SparseDiaImpl:
668 """Kronecker product of two SparseDiaImpl matrices → SparseDiaImpl.
670 For operands A (n_A × n_A) and B (m × m), the output has dimension
671 (n_A*m) × (n_A*m). Each diagonal pair (kA, kB) contributes to the
672 output diagonal at offset ``kout = kA * m + kB``.
674 Key insight — the full output diagonal can be constructed without any
675 scatter or dynamic indexing::
677 out_diag_arr[s] = left_padded[ s // m ] * right_padded[ s % m ]
679 which equals::
681 jnp.repeat(left_padded, m, axis=-1) * jnp.tile(right_padded, n_A)
683 Complexity: O(d_A * d_B * N) where N = n_A * m.
685 Args:
686 a: Left SparseDiaImpl of shape (*batch, n_A, n_A).
687 b: Right SparseDiaImpl of shape (*batch, m, m).
689 Returns:
690 SparseDiaImpl of shape (*batch, N, N).
691 """
692 n_A = a._diags.shape[-1]
693 m = b._diags.shape[-1]
695 # N = n_A * m
696 # batch_shape = jnp.broadcast_shapes(a._diags.shape[:-2], b._diags.shape[:-2])
698 # Accumulate contributions per output offset
699 out_accum: dict[int, Array] = {}
701 for li, kA in enumerate(a._offsets):
702 for ri, kB in enumerate(b._offsets):
703 kout = kA * m + kB
704 # Full output diagonal (length N) via repeat/tile — fully vectorised
705 left_rep = jnp.repeat(a._diags[..., li, :], m, axis=-1) # (*batch, N)
706 right_tiled = jnp.tile(b._diags[..., ri, :], n_A) # (*batch, N)
707 contrib = left_rep * right_tiled
708 if kout in out_accum:
709 out_accum[kout] = out_accum[kout] + contrib
710 else:
711 out_accum[kout] = contrib
713 out_offsets = tuple(sorted(out_accum.keys()))
714 out_diags = jnp.stack([out_accum[k] for k in out_offsets], axis=-2)
715 return SparseDiaImpl._make(out_offsets, out_diags)
718# ---------------------------------------------------------------------------
719# Register with the QarrayImplType enum
720# ---------------------------------------------------------------------------
722QarrayImplType.register(SparseDiaImpl, QarrayImplType.SPARSE_DIA)
725# ---------------------------------------------------------------------------
726# Public exports
727# ---------------------------------------------------------------------------
729__all__ = [
730 "SparseDiaData",
731 "SparseDiaImpl",
732]