core

Quantum Tooling

CustomProgressMeter

Bases: TqdmProgressMeter

JAXQuantum's default Diffrax progress bar.

Source code in jaxquantum/core/solvers.py

class CustomProgressMeter(diffrax.TqdmProgressMeter):
    """JAXQuantum's default Diffrax progress bar."""

    @staticmethod
    def _init_bar() -> tqdm.tqdm:
        bar_format = (
            "{desc}: {percentage:3.0f}% |{bar}| "
            "[{elapsed}<{remaining}, {rate_fmt}{postfix}]"
        )
        return tqdm.tqdm(
            total=100, bar_format=bar_format, unit="%", colour="MAGENTA", ascii="░▒█"
        )

DenseImpl

Bases: QarrayImpl

Dense implementation using JAX dense arrays.

Attributes:

Name	Type	Description
_data	Array	The underlying jnp.ndarray.

Source code in jaxquantum/core/qarray.py

@struct.dataclass
class DenseImpl(QarrayImpl):
    """Dense implementation using JAX dense arrays.

    Attributes:
        _data: The underlying ``jnp.ndarray``.
    """

    _data: Array

    PROMOTION_ORDER = 2  # noqa: RUF012 — not a struct field; no annotation intentional

    @classmethod
    def _make(cls, data) -> "DenseImpl":
        """Construct a ``DenseImpl``, applying the configured default sharding.

        All internal construction sites route through this so that every
        Qarray (including intermediates produced by ``matmul``, ``kron``,
        etc.) satisfies the user's sharding invariant.
        """
        return cls(_data=_maybe_shard(data))

    @classmethod
    def from_data(cls, data) -> "DenseImpl":
        """Wrap *data* in a new ``DenseImpl``.

        Args:
            data: Array-like input data.

        Returns:
            A ``DenseImpl`` wrapping ``robust_asarray(data)``.
        """
        return cls._make(robust_asarray(data))

    def get_data(self) -> Array:
        """Return the underlying dense array."""
        return self._data

    def matmul(self, other: QarrayImpl) -> QarrayImpl:
        """Matrix multiply ``self @ other``, coercing types as needed.

        Args:
            other: Right-hand operand.

        Returns:
            A ``DenseImpl`` containing the matrix product.
        """
        a, b = self._coerce(other)
        if a is not self:
            return a.matmul(b)
        return DenseImpl._make(self._data @ b._data)

    def add(self, other: QarrayImpl) -> QarrayImpl:
        """Element-wise addition ``self + other``, coercing types as needed.

        Args:
            other: Right-hand operand.

        Returns:
            A ``DenseImpl`` containing the sum.
        """
        a, b = self._coerce(other)
        if a is not self:
            return a.add(b)
        return DenseImpl._make(self._data + b._data)

    def sub(self, other: QarrayImpl) -> QarrayImpl:
        """Element-wise subtraction ``self - other``, coercing types as needed.

        Args:
            other: Right-hand operand.

        Returns:
            A ``DenseImpl`` containing the difference.
        """
        a, b = self._coerce(other)
        if a is not self:
            return a.sub(b)
        return DenseImpl._make(self._data - b._data)

    def mul(self, scalar) -> QarrayImpl:
        """Scalar multiplication.

        Args:
            scalar: Scalar value.

        Returns:
            A ``DenseImpl`` with each element multiplied by *scalar*.
        """
        return DenseImpl._make(scalar * self._data)

    def dag(self) -> QarrayImpl:
        """Conjugate transpose.

        Returns:
            A ``DenseImpl`` containing the conjugate transpose.
        """
        return DenseImpl._make(jnp.moveaxis(jnp.conj(self._data), -1, -2))

    def to_dense(self) -> "DenseImpl":
        """Return self (already dense).

        Returns:
            This ``DenseImpl`` instance unchanged.
        """
        return self

    def to_sparse_bcoo(self) -> "SparseBCOOImpl":
        """Convert to a ``SparseBCOOImpl`` via ``BCOO.fromdense``.

        Returns:
            A ``SparseBCOOImpl`` wrapping a BCOO conversion of this array.
        """
        from jaxquantum.core.sparse_bcoo import SparseBCOOImpl
        return SparseBCOOImpl(sparse.BCOO.fromdense(self._data))

    def shape(self) -> tuple:
        """Shape of the underlying dense array.

        Returns:
            Tuple of dimension sizes.
        """
        return self._data.shape

    def dtype(self):
        """Data type of the underlying dense array.

        Returns:
            The dtype of ``_data``.
        """
        return self._data.dtype

    def frobenius_norm(self) -> float:
        """Compute the Frobenius norm.

        Returns:
            The Frobenius norm as a scalar.
        """
        return jnp.sqrt(jnp.sum(jnp.abs(self._data) ** 2))

    def real(self) -> QarrayImpl:
        """Element-wise real part.

        Returns:
            A ``DenseImpl`` containing the real parts.
        """
        return DenseImpl._make(jnp.real(self._data))

    def imag(self) -> QarrayImpl:
        """Element-wise imaginary part.

        Returns:
            A ``DenseImpl`` containing the imaginary parts.
        """
        return DenseImpl._make(jnp.imag(self._data))

    def conj(self) -> QarrayImpl:
        """Element-wise complex conjugate.

        Returns:
            A ``DenseImpl`` containing the complex-conjugated values.
        """
        return DenseImpl._make(jnp.conj(self._data))

    def __deepcopy__(self, memo=None):
        return DenseImpl._make(deepcopy(self._data, memo))

    def tidy_up(self, atol):
        """Zero out real/imaginary parts whose magnitude is below *atol*.

        Args:
            atol: Absolute tolerance threshold.

        Returns:
            A new ``DenseImpl`` with small values zeroed.
        """
        data = self._data
        data_re = jnp.real(data)
        data_im = jnp.imag(data)
        data_re_mask = jnp.abs(data_re) > atol
        data_im_mask = jnp.abs(data_im) > atol
        data_new = data_re * data_re_mask + 1j * data_im * data_im_mask

        return DenseImpl._make(data_new)

    def kron(self, other: "QarrayImpl") -> "QarrayImpl":
        """Kronecker product using ``jnp.kron``.

        Args:
            other: Right-hand operand.

        Returns:
            A ``DenseImpl`` containing the Kronecker product.
        """
        a, b = self._coerce(other)
        if a is not self:
            return a.kron(b)
        return DenseImpl._make(jnp.kron(self._data, b._data))

    @classmethod
    def _eye_data(cls, n: int, dtype=None):
        """Create an ``n x n`` identity matrix as a dense JAX array.

        Args:
            n: Matrix size.
            dtype: Optional data type.

        Returns:
            A ``jnp.ndarray`` identity matrix of shape ``(n, n)``.
        """
        return jnp.eye(n, dtype=dtype)

    @classmethod
    def can_handle_data(cls, arr) -> bool:
        """Return True for any non-BCOO, non-SparseDIA array.

        ``SparseDiaData`` objects carry a ``_is_sparse_dia`` marker so we can
        exclude them without a direct type import (which would be circular).

        Args:
            arr: Raw array.

        Returns:
            True when *arr* is a plain dense array (not BCOO, not SparseDiaData).
        """
        return not isinstance(arr, sparse.BCOO) and not getattr(arr, "_is_sparse_dia", False)

    @classmethod
    def dag_data(cls, arr) -> Array:
        """Conjugate transpose for dense arrays.

        Swaps the last two axes via :func:`jnp.moveaxis` and conjugates all
        elements.  For 1-D inputs only conjugation is applied.

        Args:
            arr: Dense array.

        Returns:
            Conjugate transpose with the last two axes swapped.
        """
        if len(arr.shape) == 1:
            return jnp.conj(arr)
        return jnp.moveaxis(jnp.conj(arr), -1, -2)

add(other)

Element-wise addition self + other, coercing types as needed.

Parameters:

Name	Type	Description	Default
other	QarrayImpl	Right-hand operand.	required

Returns:

Type	Description
QarrayImpl	A DenseImpl containing the sum.

Source code in jaxquantum/core/qarray.py

def add(self, other: QarrayImpl) -> QarrayImpl:
    """Element-wise addition ``self + other``, coercing types as needed.

    Args:
        other: Right-hand operand.

    Returns:
        A ``DenseImpl`` containing the sum.
    """
    a, b = self._coerce(other)
    if a is not self:
        return a.add(b)
    return DenseImpl._make(self._data + b._data)

can_handle_data(arr) classmethod

Return True for any non-BCOO, non-SparseDIA array.

SparseDiaData objects carry a _is_sparse_dia marker so we can exclude them without a direct type import (which would be circular).

Parameters:

Name	Type	Description	Default
arr		Raw array.	required

Returns:

Type	Description
bool	True when arr is a plain dense array (not BCOO, not SparseDiaData).

Source code in jaxquantum/core/qarray.py

@classmethod
def can_handle_data(cls, arr) -> bool:
    """Return True for any non-BCOO, non-SparseDIA array.

    ``SparseDiaData`` objects carry a ``_is_sparse_dia`` marker so we can
    exclude them without a direct type import (which would be circular).

    Args:
        arr: Raw array.

    Returns:
        True when *arr* is a plain dense array (not BCOO, not SparseDiaData).
    """
    return not isinstance(arr, sparse.BCOO) and not getattr(arr, "_is_sparse_dia", False)

conj()

Element-wise complex conjugate.

Returns:

Type	Description
QarrayImpl	A DenseImpl containing the complex-conjugated values.

Source code in jaxquantum/core/qarray.py

def conj(self) -> QarrayImpl:
    """Element-wise complex conjugate.

    Returns:
        A ``DenseImpl`` containing the complex-conjugated values.
    """
    return DenseImpl._make(jnp.conj(self._data))

dag()

Conjugate transpose.

Returns:

Type	Description
QarrayImpl	A DenseImpl containing the conjugate transpose.

Source code in jaxquantum/core/qarray.py

def dag(self) -> QarrayImpl:
    """Conjugate transpose.

    Returns:
        A ``DenseImpl`` containing the conjugate transpose.
    """
    return DenseImpl._make(jnp.moveaxis(jnp.conj(self._data), -1, -2))

dag_data(arr) classmethod

Conjugate transpose for dense arrays.

Swaps the last two axes via :func:jnp.moveaxis and conjugates all elements. For 1-D inputs only conjugation is applied.

Parameters:

Name	Type	Description	Default
arr		Dense array.	required

Returns:

Type	Description
Array	Conjugate transpose with the last two axes swapped.

Source code in jaxquantum/core/qarray.py

@classmethod
def dag_data(cls, arr) -> Array:
    """Conjugate transpose for dense arrays.

    Swaps the last two axes via :func:`jnp.moveaxis` and conjugates all
    elements.  For 1-D inputs only conjugation is applied.

    Args:
        arr: Dense array.

    Returns:
        Conjugate transpose with the last two axes swapped.
    """
    if len(arr.shape) == 1:
        return jnp.conj(arr)
    return jnp.moveaxis(jnp.conj(arr), -1, -2)

dtype()

Data type of the underlying dense array.

Returns:

Type	Description
	The dtype of _data.

Source code in jaxquantum/core/qarray.py

def dtype(self):
    """Data type of the underlying dense array.

    Returns:
        The dtype of ``_data``.
    """
    return self._data.dtype

frobenius_norm()

Compute the Frobenius norm.

Returns:

Type	Description
float	The Frobenius norm as a scalar.

Source code in jaxquantum/core/qarray.py

def frobenius_norm(self) -> float:
    """Compute the Frobenius norm.

    Returns:
        The Frobenius norm as a scalar.
    """
    return jnp.sqrt(jnp.sum(jnp.abs(self._data) ** 2))

from_data(data) classmethod

Wrap data in a new DenseImpl.

Parameters:

Name	Type	Description	Default
data		Array-like input data.	required

Returns:

Type	Description
'DenseImpl'	A DenseImpl wrapping robust_asarray(data).

Source code in jaxquantum/core/qarray.py

@classmethod
def from_data(cls, data) -> "DenseImpl":
    """Wrap *data* in a new ``DenseImpl``.

    Args:
        data: Array-like input data.

    Returns:
        A ``DenseImpl`` wrapping ``robust_asarray(data)``.
    """
    return cls._make(robust_asarray(data))

get_data()

Return the underlying dense array.

Source code in jaxquantum/core/qarray.py

def get_data(self) -> Array:
    """Return the underlying dense array."""
    return self._data

imag()

Element-wise imaginary part.

Returns:

Type	Description
QarrayImpl	A DenseImpl containing the imaginary parts.

Source code in jaxquantum/core/qarray.py

def imag(self) -> QarrayImpl:
    """Element-wise imaginary part.

    Returns:
        A ``DenseImpl`` containing the imaginary parts.
    """
    return DenseImpl._make(jnp.imag(self._data))

kron(other)

Kronecker product using jnp.kron.

Parameters:

Name	Type	Description	Default
other	'QarrayImpl'	Right-hand operand.	required

Returns:

Type	Description
'QarrayImpl'	A DenseImpl containing the Kronecker product.

Source code in jaxquantum/core/qarray.py

def kron(self, other: "QarrayImpl") -> "QarrayImpl":
    """Kronecker product using ``jnp.kron``.

    Args:
        other: Right-hand operand.

    Returns:
        A ``DenseImpl`` containing the Kronecker product.
    """
    a, b = self._coerce(other)
    if a is not self:
        return a.kron(b)
    return DenseImpl._make(jnp.kron(self._data, b._data))

matmul(other)

Matrix multiply self @ other, coercing types as needed.

Parameters:

Name	Type	Description	Default
other	QarrayImpl	Right-hand operand.	required

Returns:

Type	Description
QarrayImpl	A DenseImpl containing the matrix product.

Source code in jaxquantum/core/qarray.py

def matmul(self, other: QarrayImpl) -> QarrayImpl:
    """Matrix multiply ``self @ other``, coercing types as needed.

    Args:
        other: Right-hand operand.

    Returns:
        A ``DenseImpl`` containing the matrix product.
    """
    a, b = self._coerce(other)
    if a is not self:
        return a.matmul(b)
    return DenseImpl._make(self._data @ b._data)

mul(scalar)

Scalar multiplication.

Parameters:

Name	Type	Description	Default
scalar		Scalar value.	required

Returns:

Type	Description
QarrayImpl	A DenseImpl with each element multiplied by scalar.

Source code in jaxquantum/core/qarray.py

def mul(self, scalar) -> QarrayImpl:
    """Scalar multiplication.

    Args:
        scalar: Scalar value.

    Returns:
        A ``DenseImpl`` with each element multiplied by *scalar*.
    """
    return DenseImpl._make(scalar * self._data)

real()

Element-wise real part.

Returns:

Type	Description
QarrayImpl	A DenseImpl containing the real parts.

Source code in jaxquantum/core/qarray.py

def real(self) -> QarrayImpl:
    """Element-wise real part.

    Returns:
        A ``DenseImpl`` containing the real parts.
    """
    return DenseImpl._make(jnp.real(self._data))

shape()

Shape of the underlying dense array.

Returns:

Type	Description
tuple	Tuple of dimension sizes.

Source code in jaxquantum/core/qarray.py

def shape(self) -> tuple:
    """Shape of the underlying dense array.

    Returns:
        Tuple of dimension sizes.
    """
    return self._data.shape

sub(other)

Element-wise subtraction self - other, coercing types as needed.

Parameters:

Name	Type	Description	Default
other	QarrayImpl	Right-hand operand.	required

Returns:

Type	Description
QarrayImpl	A DenseImpl containing the difference.

Source code in jaxquantum/core/qarray.py

def sub(self, other: QarrayImpl) -> QarrayImpl:
    """Element-wise subtraction ``self - other``, coercing types as needed.

    Args:
        other: Right-hand operand.

    Returns:
        A ``DenseImpl`` containing the difference.
    """
    a, b = self._coerce(other)
    if a is not self:
        return a.sub(b)
    return DenseImpl._make(self._data - b._data)

tidy_up(atol)

Zero out real/imaginary parts whose magnitude is below atol.

Parameters:

Name	Type	Description	Default
atol		Absolute tolerance threshold.	required

Returns:

Type	Description
	A new DenseImpl with small values zeroed.

Source code in jaxquantum/core/qarray.py

def tidy_up(self, atol):
    """Zero out real/imaginary parts whose magnitude is below *atol*.

    Args:
        atol: Absolute tolerance threshold.

    Returns:
        A new ``DenseImpl`` with small values zeroed.
    """
    data = self._data
    data_re = jnp.real(data)
    data_im = jnp.imag(data)
    data_re_mask = jnp.abs(data_re) > atol
    data_im_mask = jnp.abs(data_im) > atol
    data_new = data_re * data_re_mask + 1j * data_im * data_im_mask

    return DenseImpl._make(data_new)

to_dense()

Return self (already dense).

Returns:

Type	Description
'DenseImpl'	This DenseImpl instance unchanged.

Source code in jaxquantum/core/qarray.py

def to_dense(self) -> "DenseImpl":
    """Return self (already dense).

    Returns:
        This ``DenseImpl`` instance unchanged.
    """
    return self

to_sparse_bcoo()

Convert to a SparseBCOOImpl via BCOO.fromdense.

Returns:

Type	Description
'SparseBCOOImpl'	A SparseBCOOImpl wrapping a BCOO conversion of this array.

Source code in jaxquantum/core/qarray.py

def to_sparse_bcoo(self) -> "SparseBCOOImpl":
    """Convert to a ``SparseBCOOImpl`` via ``BCOO.fromdense``.

    Returns:
        A ``SparseBCOOImpl`` wrapping a BCOO conversion of this array.
    """
    from jaxquantum.core.sparse_bcoo import SparseBCOOImpl
    return SparseBCOOImpl(sparse.BCOO.fromdense(self._data))

Qarray

Bases: Generic[ImplT]

Quantum array with a pluggable storage backend.

Qarray wraps a QarrayImpl together with quantum-mechanical dimension metadata (_qdims) and optional batch dimensions (_bdims). The default backend is dense (DenseImpl); pass implementation="sparse_bcoo" (or QarrayImplType.SPARSE_BCOO) to store data as a JAX BCOO sparse array.

Attributes:

Name	Type	Description
_impl	ImplT	The storage backend holding the raw data.
_qdims	Qdims	Quantum dimension metadata (bra/ket structure, Hilbert space sizes).
_bdims	tuple[int]	Tuple of batch dimension sizes (empty tuple = non-batched).

Example

import jaxquantum as jqt a = jqt.destroy(10, implementation="sparse_bcoo") a.is_sparse_bcoo True

Source code in jaxquantum/core/qarray.py

@struct.dataclass
class Qarray(Generic[ImplT]):
    """Quantum array with a pluggable storage backend.

    ``Qarray`` wraps a ``QarrayImpl`` together with quantum-mechanical
    dimension metadata (``_qdims``) and optional batch dimensions
    (``_bdims``).  The default backend is dense (``DenseImpl``); pass
    ``implementation="sparse_bcoo"`` (or ``QarrayImplType.SPARSE_BCOO``) to
    store data as a JAX BCOO sparse array.

    Attributes:
        _impl: The storage backend holding the raw data.
        _qdims: Quantum dimension metadata (bra/ket structure, Hilbert space
            sizes).
        _bdims: Tuple of batch dimension sizes (empty tuple = non-batched).

    Example:
        >>> import jaxquantum as jqt
        >>> a = jqt.destroy(10, implementation="sparse_bcoo")
        >>> a.is_sparse_bcoo
        True
    """

    _impl: ImplT
    _qdims: Qdims = struct.field(pytree_node=False)
    _bdims: tuple[int] = struct.field(pytree_node=False)

    # Initialization ----
    @classmethod
    def _from_impl(cls, impl, qdims, bdims=None):
        """Wrap a canonical internal result without repeating public validation."""
        if bdims is None:
            shape = impl.shape()
            bdims = shape[:-2] if qdims.qtype == Qtypes.oper else shape[:-1]
        return cls(impl, qdims, tuple(bdims))

    @classmethod
    @overload
    def create(cls, data, dims=None, bdims=None, qtype=None, implementation: Literal[QarrayImplType.DENSE] = QarrayImplType.DENSE) -> "Qarray[DenseImpl]":
        ...

    @classmethod
    @overload
    def create(cls, data, dims=None, bdims=None, qtype=None, implementation: Literal[QarrayImplType.SPARSE_BCOO] = ...) -> "Qarray[SparseBCOOImpl]":
        ...

    @classmethod
    @overload
    def create(cls, data, dims=None, bdims=None, qtype=None, implementation=...) -> "Qarray[DenseImpl]":
        ...

    @classmethod
    def create(cls, data, dims=None, bdims=None, qtype=None, implementation=QarrayImplType.DENSE):
        """Create a ``Qarray`` from raw data.

        Handles shape normalisation, dimension inference, and tidying of small
        values.

        State vectors are stored with their Hilbert space on a **single trailing
        axis** — a ket/bra of dimension ``N`` has data shape ``bdims + (N,)``
        (never ``(N,1)`` / ``(1,N)``).  Operators keep the last two axes:
        ``bdims + (M, N)``.  The ket/bra/oper distinction lives in ``_qdims``,
        not in the data shape.

        Legacy ``(N,1)`` / ``(1,N)`` inputs are still accepted and are squeezed
        to ``(N,)`` on the way in.  Because a 1‑D ``(N,)`` array (or a square
        ``(N,N)`` batch of vectors) is shape-ambiguous, pass ``qtype`` to be
        explicit.

        Args:
            data: Input data array (dense array-like or ``sparse.BCOO``).
            dims: Quantum dimensions as ``((row_dims...), (col_dims...))`` or, as
                a convenience, a flat list of Hilbert space sizes.  Inferred from
                *data* shape when ``None``.
            bdims: Tuple of batch dimension sizes.  Inferred from the leading
                dimensions of *data* when ``None``.
            qtype: Optional quantum type — ``"ket"`` / ``"bra"`` / ``"oper"`` (or
                a :class:`Qtypes` member).  When given it is authoritative and
                resolves shape ambiguity (e.g. a bare ``(N,)`` array with
                ``qtype="ket"`` behaves like a legacy ``(N,1)`` input).
            implementation: Storage backend — ``QarrayImplType.DENSE``
                (default) or ``QarrayImplType.SPARSE_BCOO``, or the equivalent
                string ``"dense"`` / ``"sparse_bcoo"``.

        Returns:
            A new ``Qarray`` backed by the requested implementation.
        """
        data = robust_asarray(data)

        # Whether the caller explicitly requested a qtype (vs. it being derived
        # from a full ``dims`` tuple). An explicit qtype is authoritative and
        # mismatches with the data raise a clear error.
        qtype_explicit = qtype is not None
        if qtype is not None and not isinstance(qtype, Qtypes):
            # Accepts "ket" / "bra" / "oper" (case-sensitive enum value).
            qtype = Qtypes.from_str(qtype)

        # Empty array (e.g. an empty list of operators) — keep legacy behaviour.
        if len(data.shape) == 1 and data.shape[0] == 0:
            dims = ((), ()) if dims is None else (tuple(dims[0]), tuple(dims[1]))
            bdims = (0,) if bdims is None else tuple(bdims)
            check_dims(dims, bdims, data.shape)
            qdims = Qdims(dims)
            impl_class = QarrayImplType(implementation).get_impl_class()
            impl = impl_class.from_data(data)
            impl = impl.tidy_up(SETTINGS["auto_tidyup_atol"])
            return cls(impl, qdims, bdims)

        # Resolve full quantum dims + batch dims.
        dims_is_full = dims is not None and isinstance(dims[0], (list, tuple))
        if qtype is None and dims_is_full:
            qtype = Qtypes.from_dims((tuple(dims[0]), tuple(dims[1])))

        if qtype is not None:
            # Explicit (or dims-derived) qtype: unambiguous layout.
            data, dims, bdims = cls._qtype_layout(data, dims, bdims, qtype)
        else:
            # Auto: legacy shape heuristics (1‑D → ket, square → oper, etc.).
            data, dims, bdims = cls._legacy_layout(data, dims, bdims)

        qdims = Qdims(dims)

        # Store vectors as bdims + (N,); operators as bdims + (M, N).
        if qdims.qtype == Qtypes.oper:
            space = (prod(dims[0]), prod(dims[1]))
        else:
            space = (prod(dims[0]) * prod(dims[1]),)

        target_shape = tuple(bdims) + space

        # An explicitly requested qtype is authoritative: error out if the data
        # cannot be interpreted as that type rather than silently coercing it.
        if qtype_explicit:
            if qdims.qtype != qtype:
                raise ValueError(
                    f"Requested qtype={qtype} is incompatible with the resolved "
                    f"dims {dims} (which is a {qdims.qtype})."
                )
            if prod(data.shape) != prod(target_shape):
                raise ValueError(
                    f"Data of shape {tuple(data.shape)} is incompatible with "
                    f"qtype={qtype} and dims={dims}: cannot reshape to "
                    f"{target_shape}."
                )

        # Reshape to the canonical stored shape only when needed. Use varargs so
        # this works for dense arrays, BCOO, and SparseDiaData alike.
        if tuple(data.shape) != target_shape:
            data = data.reshape(*target_shape)

        check_dims(dims, bdims, data.shape)

        # NOTE: Constantly tidying up on Qarray creation might be a bit overkill.
        # It increases the compilation time, but only very slightly
        # increased the runtime of the jit compiled function.
        # We could instead use this tidy up where we think we need it.

        impl_class = QarrayImplType(implementation).get_impl_class()
        impl = impl_class.from_data(data)
        impl = impl.tidy_up(SETTINGS["auto_tidyup_atol"])

        return cls(impl, qdims, bdims)

    @staticmethod
    def _legacy_layout(data, dims, bdims):
        """Auto-resolve dims/bdims from data shape (no explicit ``qtype``).

        Mirrors the historical shape heuristics: a 1‑D array is a ket, a square
        2‑D array is an operator, a non-square 2‑D array is a batch of kets, and
        legacy ``(N,1)`` / ``(1,N)`` carry the orientation.  Returns
        ``(data, dims, bdims)`` where *data* may still carry a trailing singleton
        for vectors — the caller collapses it to ``(N,)``.
        """
        if len(data.shape) == 1 and data.shape[0] > 0:
            data = data.reshape(data.shape[0], 1)

        if (
            len(data.shape) >= 2
            and data.shape[-2] != data.shape[-1]
            and not (data.shape[-2] == 1 or data.shape[-1] == 1)
        ):
            data = data.reshape(*data.shape[:-1], data.shape[-1], 1)

        if bdims is not None and len(data.shape) - len(bdims) == 1:
            data = data.reshape(*data.shape[:-1], data.shape[-1], 1)

        if bdims is None:
            bdims = tuple(data.shape[:-2])

        if dims is None:
            dims = ((data.shape[-2],), (data.shape[-1],))

        if not isinstance(dims[0], (list, tuple)):
            # Only the Hilbert space dimensions were sent in.
            if data.shape[-1] == 1:
                dims = (tuple(dims), tuple([1 for _ in dims]))
            elif data.shape[-2] == 1:
                dims = (tuple([1 for _ in dims]), tuple(dims))
            else:
                dims = (tuple(dims), tuple(dims))
        else:
            dims = (tuple(dims[0]), tuple(dims[1]))

        return data, dims, tuple(bdims)

    @staticmethod
    def _qtype_layout(data, dims, bdims, qtype):
        """Resolve dims/bdims given an explicit ``qtype`` (no shape guessing).

        Accepts modern ``(...,N)`` vectors as well as legacy ``(N,1)`` / ``(1,N)``
        (and batched variants) and returns ``(data, dims_full, bdims)``.  The
        caller reshapes *data* to the canonical stored shape.
        """
        shape = tuple(data.shape)
        dims_is_full = dims is not None and isinstance(dims[0], (list, tuple))

        if dims_is_full:
            dims_full = (tuple(dims[0]), tuple(dims[1]))
        elif dims is not None:
            sl = tuple(dims)
            if qtype == Qtypes.ket:
                dims_full = (sl, tuple(1 for _ in sl))
            elif qtype == Qtypes.bra:
                dims_full = (tuple(1 for _ in sl), sl)
            else:
                dims_full = (sl, sl)
        else:
            dims_full = None  # inferred from shape below

        if qtype == Qtypes.oper:
            if dims_full is None:
                if len(shape) < 2:
                    raise ValueError(
                        f"qtype='oper' needs 2-D data or explicit dims; "
                        f"got shape {shape}."
                    )
                dims_full = ((shape[-2],), (shape[-1],))
            if bdims is None:
                bdims = shape[:-2]
            return data, dims_full, tuple(bdims)

        # Vector (ket / bra).
        if bdims is not None:
            bdims = tuple(bdims)
            if dims_full is None:
                n = 1
                for d in shape[len(bdims):]:
                    n *= d
                dims_full = ((n,), (1,)) if qtype == Qtypes.ket else ((1,), (n,))
            return data, dims_full, bdims

        # bdims unknown: peel a legacy orientation singleton, then the last axis
        # is the space axis and the leading axes are batch dims.
        s = list(shape)
        if qtype == Qtypes.ket and len(s) >= 2 and s[-1] == 1:
            s = s[:-1]
        elif qtype == Qtypes.bra and len(s) >= 2 and s[-2] == 1:
            s = s[:-2] + s[-1:]
        inferred_n = s[-1] if s else 1
        bdims = tuple(s[:-1])
        if dims_full is None:
            dims_full = ((inferred_n,), (1,)) if qtype == Qtypes.ket else ((1,), (inferred_n,))
        return data, dims_full, bdims

    @classmethod
    @overload
    def from_sparse_bcoo(cls, data, dims=None, bdims=None) -> "Qarray[SparseBCOOImpl]":
        ...

    @classmethod
    def from_sparse_bcoo(cls, data, dims=None, bdims=None):
        """Create a ``Qarray`` directly from a sparse BCOO array without densifying.

        Args:
            data: A ``sparse.BCOO`` or array-like to store as sparse BCOO.
            dims: Quantum dimensions.  Inferred when ``None``.
            bdims: Batch dimensions.  Inferred when ``None``.

        Returns:
            A ``Qarray[SparseBCOOImpl]``.
        """
        return cls.create(data, dims=dims, bdims=bdims, implementation=QarrayImplType.SPARSE_BCOO)

    @classmethod
    def from_sparse_dia(cls, data, dims=None, bdims=None) -> "Qarray":
        """Create a SparseDIA-backed ``Qarray``.

        Accepts either a dense array-like (diagonals are auto-detected) or a
        :class:`~jaxquantum.core.sparse_dia.SparseDiaData` container.

        Args:
            data: Dense array of shape (*batch, n, n) or a ``SparseDiaData``.
            dims: Quantum dimensions ``((row_dims,), (col_dims,))``.
            bdims: Batch dimension sizes.

        Returns:
            A ``Qarray`` backed by ``SparseDiaImpl``.
        """
        return cls.create(data, dims=dims, bdims=bdims, implementation=QarrayImplType.SPARSE_DIA)

    @classmethod
    @overload
    def from_list(cls, qarr_list: List["Qarray[DenseImpl]"], qtype=None) -> "Qarray[DenseImpl]":
        ...

    @classmethod
    @overload
    def from_list(cls, qarr_list: List["Qarray[SparseBCOOImpl]"], qtype=None) -> "Qarray[SparseBCOOImpl]":
        ...

    @classmethod
    def from_list(cls, qarr_list: List[Qarray], qtype=None) -> Qarray:
        """Create a batched ``Qarray`` from a list of same-shaped ``Qarray`` objects.

        The output implementation is determined by the element with the highest
        ``PROMOTION_ORDER``: if all inputs are sparse the result is sparse; if
        any input is dense (or types are mixed) all inputs are promoted to dense
        and the result is dense.

        Works for kets/bras (stacked into ``(len, *bdims, N)``) as well as
        operators, regardless of whether the elements were originally created
        from ``(N,)`` or legacy ``(N,1)`` / ``(1,N)`` arrays — they are all
        stored as ``(N,)`` vectors by the time they reach here.

        Args:
            qarr_list: List of ``Qarray`` objects with identical ``dims`` and
                ``bdims``.  May be empty.
            qtype: Optional quantum type ("ket"/"bra"/"oper" or a ``Qtypes``)
                forwarded to :meth:`create`.  Defaults to the qtype of the
                first element, which is the correct choice in all normal cases;
                pass it explicitly to override or to be defensive.

        Returns:
            A ``Qarray`` with an extra leading batch dimension of size
            ``len(qarr_list)``.

        Raises:
            ValueError: If the elements have mismatched ``dims`` or ``bdims``.
        """
        if len(qarr_list) == 0:
            dims = ((), ())
            bdims = (0,)
            return cls.create(jnp.array([]), dims=dims, bdims=bdims)

        dims = qarr_list[0].dims
        bdims = qarr_list[0].bdims

        if not all(qarr.dims == dims and qarr.bdims == bdims for qarr in qarr_list):
            raise ValueError("All Qarrays in the list must have the same dimensions.")

        if qtype is None:
            qtype = qarr_list[0].qtype

        new_bdims = (len(qarr_list),) + bdims

        # Pick the target type: highest PROMOTION_ORDER wins (dense beats sparse).
        target_impl_type = max(
            (q.impl_type for q in qarr_list),
            key=lambda t: t.get_impl_class().PROMOTION_ORDER,
        )

        if target_impl_type == QarrayImplType.SPARSE_DIA:
            # All inputs are SparseDIA — batch without densifying.
            # Compute union of offsets across all operators, then remap each
            # operator's _diags rows into the union shape and stack.
            from jaxquantum.core.sparse_dia import SparseDiaData  # lazy to avoid circular
            union_offsets = tuple(sorted(
                set().union(*[set(q._impl._offsets) for q in qarr_list])
            ))
            union_idx = {k: i for i, k in enumerate(union_offsets)}
            n = qarr_list[0]._impl._diags.shape[-1]
            dtype = jnp.result_type(*[q._impl._diags.dtype for q in qarr_list])
            remapped = []
            for q in qarr_list:
                row = jnp.zeros((len(union_offsets), n), dtype=dtype)
                for i_src, k in enumerate(q._impl._offsets):
                    row = row.at[union_idx[k], :].set(q._impl._diags[i_src, :])
                remapped.append(row)
            stacked = jnp.stack(remapped, axis=0)  # (n_ops, n_union_diags, N)
            raw = SparseDiaData(offsets=union_offsets, diags=stacked)
            return cls.create(raw, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.SPARSE_DIA)

        if target_impl_type == QarrayImplType.SPARSE_BCOO:
            # All inputs are sparse BCOO — stack via dense intermediates then re-sparsify.
            data = jnp.array([q.data.todense() for q in qarr_list])
            return cls.create(data, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.SPARSE_BCOO)

        # Target is dense: promote any sparse inputs before stacking.
        data = jnp.array([q.to_dense().data for q in qarr_list])
        return cls.create(data, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.DENSE)

    @classmethod
    @overload
    def from_array(cls, qarr_arr: "Qarray[DenseImpl]") -> "Qarray[DenseImpl]":
        ...

    @classmethod
    @overload
    def from_array(cls, qarr_arr: "Qarray[SparseBCOOImpl]") -> "Qarray[SparseBCOOImpl]":
        ...

    @classmethod
    def from_array(cls, qarr_arr) -> Qarray:
        """Create a ``Qarray`` from a (possibly nested) list of ``Qarray`` objects.

        Args:
            qarr_arr: A ``Qarray`` (returned as-is) or a nested list of
                ``Qarray`` objects.

        Returns:
            A ``Qarray`` with batch dimensions matching the nesting structure
            of *qarr_arr*.
        """
        if isinstance(qarr_arr, Qarray):
            return qarr_arr

        bdims = ()
        lvl = qarr_arr
        while not isinstance(lvl, Qarray):
            bdims = bdims + (len(lvl),)
            if len(lvl) > 0:
                lvl = lvl[0]
            else:
                break

        def flat(lis):
            flatList = []
            for element in lis:
                if type(element) is list:
                    flatList += flat(element)
                else:
                    flatList.append(element)
            return flatList

        qarr_list = flat(qarr_arr)
        qarr = cls.from_list(qarr_list)
        qarr = qarr.reshape_bdims(*bdims)
        return qarr

    # Properties ----
    @property
    def qtype(self):
        """Quantum type of this array (ket, bra, or operator)."""
        return self._qdims.qtype

    @property
    def dtype(self):
        """Data type of the underlying storage array."""
        return self._impl.dtype()

    @property
    def dims(self):
        """Quantum dimensions as ``((row_dims...), (col_dims...))``."""
        return self._qdims.dims

    @property
    def bdims(self):
        """Tuple of batch dimension sizes (empty tuple = non-batched)."""
        return self._bdims

    @property
    def qdims(self):
        """The ``Qdims`` metadata object for this array."""
        return self._qdims

    @property
    def space_dims(self):
        """Hilbert space dimensions for the relevant side (ket row / bra col)."""
        if self.qtype in [Qtypes.oper, Qtypes.ket]:
            return self.dims[0]
        elif self.qtype == Qtypes.bra:
            return self.dims[1]
        else:
            # TODO: not reached for some reason
            raise ValueError("Unsupported qtype.")

    @property
    def data(self):
        """The raw underlying data (dense ``jnp.ndarray`` or ``sparse.BCOO``)."""
        return self._impl.data

    @property
    def shaped_data(self):
        """Data reshaped to ``bdims + dims[0] + dims[1]``."""
        return self.data.reshape(self.bdims + self.dims[0] + self.dims[1])

    @property
    def shape(self):
        """Shape of the underlying data array."""
        return self.data.shape

    @property
    def is_batched(self):
        """True if this array has one or more batch dimensions."""
        return len(self.bdims) > 0

    @property
    def is_sparse_bcoo(self):
        """True if the storage backend is ``SparseBCOOImpl`` (BCOO)."""
        return self._impl.impl_type == QarrayImplType.SPARSE_BCOO

    @property
    def is_dense(self):
        """True if the storage backend is ``DenseImpl``."""
        return self._impl.impl_type == QarrayImplType.DENSE

    @property
    def is_sparse_dia(self):
        """True if the storage backend is ``SparseDiaImpl``."""
        return self._impl.impl_type == QarrayImplType.SPARSE_DIA

    @property
    def impl_type(self):
        """The ``QarrayImplType`` member of the current storage backend."""
        return self._impl.impl_type

    def to_sparse_bcoo(self) -> "Qarray[SparseBCOOImpl]":
        """Return a BCOO-sparse-backed copy of this array.

        If the array is already sparse BCOO, returns self unchanged.

        Returns:
            A ``Qarray[SparseBCOOImpl]``.
        """
        if self.is_sparse_bcoo:
            return self
        new_impl = self._impl.to_sparse_bcoo()
        return Qarray(new_impl, self._qdims, self._bdims)

    def to_sparse_dia(self) -> "Qarray":
        """Return a SparseDIA-backed copy of this array.

        If the array is already SparseDIA, returns self unchanged.

        Returns:
            A ``Qarray[SparseDiaImpl]``.
        """
        if self.is_sparse_dia:
            return self
        new_impl = self._impl.to_sparse_dia()
        return Qarray(new_impl, self._qdims, self._bdims)

    def to_dense(self) -> "Qarray[DenseImpl]":
        """Return a dense-backed copy of this array.

        If the array is already dense, returns self unchanged.

        Returns:
            A ``Qarray[DenseImpl]``.
        """
        if self.is_dense:
            return self
        new_impl = self._impl.to_dense()
        return Qarray(new_impl, self._qdims, self._bdims)

    def __getitem__(self, index):
        if len(self.bdims) > 0:
            impl = type(self._impl).from_data(self.data[index])
            return Qarray._from_impl(impl, self._qdims)
        else:
            raise ValueError("Cannot index a non-batched Qarray.")

    def reshape_bdims(self, *args):
        """Reshape the batch dimensions of this ``Qarray``.

        Args:
            *args: New batch dimension sizes.

        Returns:
            A new ``Qarray`` with the requested batch shape.
        """
        new_bdims = tuple(args)

        if prod(new_bdims) == 0:
            new_shape = new_bdims
        elif self.qtype == Qtypes.oper:
            new_shape = new_bdims + (prod(self.dims[0]), -1)
        else:
            # Vectors keep a single trailing space axis (no (N,1)).
            new_shape = new_bdims + (prod(self.dims[0]) * prod(self.dims[1]),)

        impl = type(self._impl).from_data(self.data.reshape(new_shape))
        return Qarray._from_impl(
            impl,
            self._qdims,
            new_bdims,
        )

    def space_to_qdims(self, space_dims: List[int]):
        """Convert Hilbert space dimensions to full quantum dims tuple.

        Args:
            space_dims: Sequence of per-subsystem Hilbert space sizes, or a
                full ``((row_dims), (col_dims))`` tuple (returned unchanged).

        Returns:
            A ``((row_dims...), (col_dims...))`` tuple.

        Raises:
            ValueError: If ``self.qtype`` is not ket, bra, or oper.
        """
        if isinstance(space_dims[0], (list, tuple)):
            return space_dims

        if self.qtype == Qtypes.oper:
            return (tuple(space_dims), tuple(space_dims))
        elif self.qtype == Qtypes.ket:
            return (tuple(space_dims), tuple([1 for _ in range(len(space_dims))]))
        elif self.qtype == Qtypes.bra:
            return (tuple([1 for _ in range(len(space_dims))]), tuple(space_dims))
        else:
            raise ValueError("Unsupported qtype for space_to_qdims conversion.")

    def reshape_qdims(self, *args):
        """Reshape the quantum dimensions of the Qarray.

        Note that this does not take in qdims but rather the new Hilbert space
        dimensions.

        Args:
            *args: New Hilbert dimensions for the Qarray.

        Returns:
            Qarray: reshaped Qarray.
        """

        new_space_dims = tuple(args)
        current_space_dims = self.space_dims
        assert prod(new_space_dims) == prod(current_space_dims)

        new_qdims = self.space_to_qdims(new_space_dims)
        return Qarray._from_impl(self._impl, Qdims(new_qdims), self._bdims)

    def resize(self, new_shape):
        """Resize the Qarray to a new shape.

        TODO: review and maybe deprecate this method.

        Args:
            new_shape: Target shape tuple.

        Returns:
            A new ``Qarray`` with data resized via ``jnp.resize``.
        """
        dims = self.dims
        data = jnp.resize(self.data, new_shape)
        # Preserve implementation type
        implementation = self.impl_type
        return Qarray.create(
            data,
            dims=dims,
            implementation=implementation,
        )

    def __len__(self):
        """Length along the first batch dimension.

        Returns:
            Size of the leading batch dimension.

        Raises:
            ValueError: If the array is not batched.
        """
        if len(self.bdims) > 0:
            return self.data.shape[0]
        else:
            raise ValueError("Cannot get length of a non-batched Qarray.")

    def __eq__(self, other):
        if not isinstance(other, Qarray):
            raise ValueError(  # noqa: TRY004
                "Cannot calculate equality of a Qarray with a non-Qarray."
            )

        if self.dims != other.dims:
            return False

        if self.bdims != other.bdims:
            return False

        if self.is_sparse_bcoo and other.is_sparse_bcoo:
            # Fast structural path: same sparsity pattern → compare values only (no todense)
            if (self.data.indices.shape == other.data.indices.shape
                    and bool(jnp.all(self.data.indices == other.data.indices))):
                return bool(jnp.allclose(self.data.data, other.data.data))
            # Different patterns: fall back to dense comparison (unavoidable)
            return bool(jnp.all(self.data.todense() == other.data.todense()))

        # At least one dense: convert sparse side to dense for comparison
        self_data  = self.data.todense()  if hasattr(self.data,  'todense') else self.data
        other_data = other.data.todense() if hasattr(other.data, 'todense') else other.data
        return bool(jnp.all(self_data == other_data))

    def __ne__(self, other):
        return not self.__eq__(other)

    # Elementary Math ----
    def __matmul__(self, other):
        if not isinstance(other, Qarray):
            return NotImplemented

        _qdims_new = self._qdims @ other._qdims
        st, ot = self.qtype, other.qtype

        # Operator @ operator: keep the backend-native path (dense or sparse),
        # which preserves the storage implementation. Operators are 2-D-tailed,
        # so create() infers bdims/qtype from the result shape unambiguously.
        if st == Qtypes.oper and ot == Qtypes.oper:
            new_impl = self._impl.matmul(other._impl)
            return Qarray._from_impl(new_impl, _qdims_new)

        # bra @ oper  ==  dag( oper^dag @ bra^dag ).  Reusing the oper @ ket path
        # keeps a sparse operator sparse (no densification of the operator).
        if st == Qtypes.bra and ot == Qtypes.oper:
            return (other.dag() @ self.dag()).dag()

        # oper @ ket: matrix-vector product. Vectors are dense (see migration
        # scope); the operator may be sparse and acts natively on the dense
        # vector. The vector's single space axis is contracted in place — it is
        # never reshaped to (N,1) for storage.
        if st == Qtypes.oper and ot == Qtypes.ket:
            vec = other.to_dense().data  # (..., N)
            if self.is_dense:
                out = jnp.einsum("...ij,...j->...i", self.data, vec)
            else:
                # Transient column only inside the sparse kernel; never stored.
                rhs = vec[..., None]  # (..., N, 1)
                out = self._impl.matmul(DenseImpl._make(rhs)).data[..., 0]
            return Qarray._from_impl(DenseImpl._make(out), _qdims_new)

        # Vector ⊗/· vector — both operands are dense.
        a = self.to_dense().data
        b = other.to_dense().data
        if st == Qtypes.ket and ot == Qtypes.bra:
            # Outer product |a><b| → operator.
            out = jnp.einsum("...i,...j->...ij", a, b)
            return Qarray._from_impl(DenseImpl._make(out), _qdims_new)
        if st == Qtypes.bra and ot == Qtypes.ket:
            # Inner product <a|b> → 1x1 "ket" (qdims ((1,),(1,))).
            out = jnp.einsum("...i,...i->...", a, b)
            out = out.reshape(out.shape + (1,))
            return Qarray._from_impl(DenseImpl._make(out), _qdims_new)

        return NotImplemented

    def __mul__(self, other):
        if isinstance(other, Qarray):
            return self.__matmul__(other)

        other = other + 0.0j
        if not robust_isscalar(other) and len(other.shape) > 0:  # not a scalar
            # Broadcast per-batch scalars against the stored data: vectors carry
            # a single trailing space axis, operators carry two.
            extra = (1,) if self.qtype in (Qtypes.ket, Qtypes.bra) else (1, 1)
            other = other.reshape(other.shape + extra)

        new_impl = self._impl.mul(other)
        return Qarray._from_impl(new_impl, self._qdims)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __neg__(self):
        return self.__mul__(-1)

    def __truediv__(self, other):
        """Divide by a scalar.

        Args:
            other: Scalar divisor.

        Returns:
            A new ``Qarray`` with all elements divided by *other*.

        Raises:
            ValueError: If *other* is a ``Qarray``.
        """
        if isinstance(other, Qarray):
            raise ValueError("Cannot divide a Qarray by another Qarray.")  # noqa: TRY004

        return self.__mul__(1 / other)

    def _scaled_identity(self, scalar):
        return type(self._impl)._scaled_identity(
            self.data.shape[-2],
            scalar,
            dtype=self.data.dtype,
        )

    def __add__(self, other):
        if isinstance(other, Qarray):
            if self.dims != other.dims:
                msg = (
                    "Dimensions are incompatible: "
                    + repr(self.dims)
                    + " and "
                    + repr(other.dims)
                )
                raise ValueError(msg)
            new_impl = self._impl.add(other._impl)
            return Qarray._from_impl(new_impl, self._qdims)

        if isinstance(other, Number) and other == 0:
            return self.copy()

        if self.qtype == Qtypes.oper:
            return Qarray._from_impl(
                self._impl.add(self._scaled_identity(other)),
                self._qdims,
            )

        return NotImplemented

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        if isinstance(other, Qarray):
            if self.dims != other.dims:
                msg = (
                    "Dimensions are incompatible: "
                    + repr(self.dims)
                    + " and "
                    + repr(other.dims)
                )
                raise ValueError(msg)
            new_impl = self._impl.sub(other._impl)
            return Qarray._from_impl(new_impl, self._qdims)

        if isinstance(other, Number) and other == 0:
            return self.copy()

        if self.qtype == Qtypes.oper:
            return Qarray._from_impl(
                self._impl.sub(self._scaled_identity(other)),
                self._qdims,
            )

        return NotImplemented

    def __rsub__(self, other):
        return self.__neg__().__add__(other)

    def __xor__(self, other):
        if not isinstance(other, Qarray):
            return NotImplemented
        return tensor(self, other)

    def __rxor__(self, other):
        if not isinstance(other, Qarray):
            return NotImplemented
        return tensor(other, self)

    def __pow__(self, other):
        if not isinstance(other, int):
            return NotImplemented

        return powm(self, other)

    # String Representation ----
    def _str_header(self):
        """Build the one-line header string for ``__str__`` and ``__repr__``."""
        impl_type = self.impl_type.value
        out = ", ".join(
            [
                "Quantum array: dims = " + str(self.dims),
                "bdims = " + str(self.bdims),
                "shape = " + str(self.data.shape),
                "type = " + str(self.qtype),
                "impl = " + impl_type,
            ]
        )
        return out

    def __str__(self):
        return self._str_header() + "\nQarray data =\n" + str(self.data)

    @property
    def header(self):
        """One-line header string describing dimensions, shape, and backend."""
        return self._str_header()

    def __repr__(self):
        return self.__str__()

    # Utilities ----
    def copy(self, memo=None):
        """Return a deep copy of this ``Qarray``.

        Args:
            memo: Optional memo dict forwarded to ``deepcopy``.

        Returns:
            A new ``Qarray`` with independent copies of all data.
        """
        return self.__deepcopy__(memo)

    def __deepcopy__(self, memo):
        """Need to override this when defining __getattr__."""

        return Qarray(
            _impl=deepcopy(self._impl, memo=memo),
            _qdims=deepcopy(self._qdims, memo=memo),
            _bdims=deepcopy(self._bdims, memo=memo),
        )

    def __getattr__(self, method_name):
        if "__" == method_name[:2]:
            # NOTE: we return NotImplemented for binary special methods logic in python, plus things like __jax_array__
            return lambda *args, **kwargs: NotImplemented

        modules = [jnp, jnp.linalg, jsp, jsp.linalg]

        method_f = None
        for mod in modules:
            method_f = getattr(mod, method_name, None)
            if method_f is not None:
                break

        if method_f is None:
            raise NotImplementedError(
                f"Method {method_name} does not exist. No backup method found in {modules}."
            )

        def func(*args, **kwargs):
            # For operations that might not be supported in sparse, convert to dense
            if self.is_sparse_bcoo:
                dense_self = self.to_dense()
                res = method_f(dense_self.data, *args, **kwargs)
            else:
                res = method_f(self.data, *args, **kwargs)

            if getattr(res, "shape", None) is None or res.shape != self.data.shape:
                return res
            else:
                # Preserve implementation type
                return Qarray.create(res, dims=self._qdims.dims, implementation=self.impl_type)

        return func

    # Conversions / Reshaping ----
    def dag(self):
        """Conjugate transpose of this array."""
        return dag(self)

    def to_dm(self):
        """Convert a ket to a density matrix via outer product."""
        return ket2dm(self)

    def is_dm(self):
        """Return True if this array is an operator (density-matrix type)."""
        return self.qtype == Qtypes.oper

    def is_vec(self):
        """Return True if this array is a ket or bra."""
        return self.qtype == Qtypes.ket or self.qtype == Qtypes.bra

    def to_ket(self):
        """Convert a bra to a ket (no-op for kets)."""
        return to_ket(self)

    def transpose(self, *args):
        """Transpose subsystem indices."""
        return transpose(self, *args)

    def keep_only_diag_elements(self):
        """Zero out all off-diagonal elements."""
        return keep_only_diag_elements(self)

    # Math Functions ----
    def unit(self):
        """Return the normalised (unit-norm) version of this array."""
        return unit(self)

    def norm(self):
        """Compute the norm of this array."""
        return norm(self)

    def frobenius_norm(self):
        """Compute the Frobenius norm directly from the implementation.

        Returns:
            The Frobenius norm as a scalar.
        """
        return self._impl.frobenius_norm()

    def real(self):
        """Element-wise real part.

        Returns:
            A new ``Qarray`` containing the real parts of each element.
        """
        new_impl = self._impl.real()
        return Qarray._from_impl(new_impl, self._qdims)

    def imag(self):
        """Element-wise imaginary part.

        Returns:
            A new ``Qarray`` containing the imaginary parts of each element.
        """
        new_impl = self._impl.imag()
        return Qarray._from_impl(new_impl, self._qdims)

    def conj(self):
        """Element-wise complex conjugate.

        Returns:
            A new ``Qarray`` containing the complex-conjugated elements.
        """
        new_impl = self._impl.conj()
        return Qarray._from_impl(new_impl, self._qdims)

    def expm(self):
        """Matrix exponential."""
        return expm(self)

    def powm(self, n):
        """Matrix power.

        Args:
            n: Exponent (integer or float).

        Returns:
            This array raised to the *n*-th matrix power.
        """
        return powm(self, n)

    def cosm(self):
        """Matrix cosine."""
        return cosm(self)

    def sinm(self):
        """Matrix sine."""
        return sinm(self)

    def tr(self, **kwargs):
        """Full trace."""
        return tr(self, **kwargs)

    def trace(self, **kwargs):
        """Full trace (alias for :meth:`tr`)."""
        return tr(self, **kwargs)

    def ptrace(self, indx):
        """Partial trace over subsystem *indx*.

        Args:
            indx: Index of the subsystem to trace out.

        Returns:
            Reduced density matrix.
        """
        return ptrace(self, indx)

    def eigenstates(self):
        """Eigenvalues and eigenstates of this operator."""
        return eigenstates(self)

    def eigenenergies(self):
        """Eigenvalues of this operator."""
        return eigenenergies(self)

    def eigenvalues(self):
        """Eigenvalues of this operator (alias for :meth:`eigenenergies`)."""
        return eigenenergies(self)

    def collapse(self, mode="sum"):
        """Collapse batch dimensions.

        Args:
            mode: Collapse strategy — currently only ``"sum"`` is supported.

        Returns:
            A non-batched ``Qarray``.
        """
        return collapse(self, mode=mode)

bdims property

Tuple of batch dimension sizes (empty tuple = non-batched).

data property

The raw underlying data (dense jnp.ndarray or sparse.BCOO).

dims property

Quantum dimensions as ((row_dims...), (col_dims...)).

dtype property

Data type of the underlying storage array.

header property

One-line header string describing dimensions, shape, and backend.

impl_type property

The QarrayImplType member of the current storage backend.

is_batched property

True if this array has one or more batch dimensions.

is_dense property

True if the storage backend is DenseImpl.

is_sparse_bcoo property

True if the storage backend is SparseBCOOImpl (BCOO).

is_sparse_dia property

True if the storage backend is SparseDiaImpl.

qdims property

The Qdims metadata object for this array.

qtype property

Quantum type of this array (ket, bra, or operator).

shape property

Shape of the underlying data array.

shaped_data property

Data reshaped to bdims + dims[0] + dims[1].

space_dims property

Hilbert space dimensions for the relevant side (ket row / bra col).

__deepcopy__(memo)

Need to override this when defining getattr.

Source code in jaxquantum/core/qarray.py

def __deepcopy__(self, memo):
    """Need to override this when defining __getattr__."""

    return Qarray(
        _impl=deepcopy(self._impl, memo=memo),
        _qdims=deepcopy(self._qdims, memo=memo),
        _bdims=deepcopy(self._bdims, memo=memo),
    )

__len__()

Length along the first batch dimension.

Returns:

Type	Description
	Size of the leading batch dimension.

Raises:

Type	Description
ValueError	If the array is not batched.

Source code in jaxquantum/core/qarray.py

def __len__(self):
    """Length along the first batch dimension.

    Returns:
        Size of the leading batch dimension.

    Raises:
        ValueError: If the array is not batched.
    """
    if len(self.bdims) > 0:
        return self.data.shape[0]
    else:
        raise ValueError("Cannot get length of a non-batched Qarray.")

__truediv__(other)

Divide by a scalar.

Parameters:

Name	Type	Description	Default
other		Scalar divisor.	required

Returns:

Type	Description
	A new Qarray with all elements divided by other.

Raises:

Type	Description
ValueError	If other is a Qarray.

Source code in jaxquantum/core/qarray.py

def __truediv__(self, other):
    """Divide by a scalar.

    Args:
        other: Scalar divisor.

    Returns:
        A new ``Qarray`` with all elements divided by *other*.

    Raises:
        ValueError: If *other* is a ``Qarray``.
    """
    if isinstance(other, Qarray):
        raise ValueError("Cannot divide a Qarray by another Qarray.")  # noqa: TRY004

    return self.__mul__(1 / other)

collapse(mode='sum')

Collapse batch dimensions.

Parameters:

Name	Type	Description	Default
mode		Collapse strategy — currently only "sum" is supported.	'sum'

Returns:

Type	Description
	A non-batched Qarray.

Source code in jaxquantum/core/qarray.py

def collapse(self, mode="sum"):
    """Collapse batch dimensions.

    Args:
        mode: Collapse strategy — currently only ``"sum"`` is supported.

    Returns:
        A non-batched ``Qarray``.
    """
    return collapse(self, mode=mode)

conj()

Element-wise complex conjugate.

Returns:

Type	Description
	A new Qarray containing the complex-conjugated elements.

Source code in jaxquantum/core/qarray.py

def conj(self):
    """Element-wise complex conjugate.

    Returns:
        A new ``Qarray`` containing the complex-conjugated elements.
    """
    new_impl = self._impl.conj()
    return Qarray._from_impl(new_impl, self._qdims)

copy(memo=None)

Return a deep copy of this Qarray.

Parameters:

Name	Type	Description	Default
memo		Optional memo dict forwarded to deepcopy.	None

Returns:

Type	Description
	A new Qarray with independent copies of all data.

Source code in jaxquantum/core/qarray.py

def copy(self, memo=None):
    """Return a deep copy of this ``Qarray``.

    Args:
        memo: Optional memo dict forwarded to ``deepcopy``.

    Returns:
        A new ``Qarray`` with independent copies of all data.
    """
    return self.__deepcopy__(memo)

cosm()

Matrix cosine.

Source code in jaxquantum/core/qarray.py

def cosm(self):
    """Matrix cosine."""
    return cosm(self)

create(data, dims=None, bdims=None, qtype=None, implementation=QarrayImplType.DENSE) classmethod

create(data, dims=None, bdims=None, qtype=None, implementation: Literal[QarrayImplType.DENSE] = QarrayImplType.DENSE) -> 'Qarray[DenseImpl]'

create(data, dims=None, bdims=None, qtype=None, implementation: Literal[QarrayImplType.SPARSE_BCOO] = ...) -> 'Qarray[SparseBCOOImpl]'

create(data, dims=None, bdims=None, qtype=None, implementation=...) -> 'Qarray[DenseImpl]'

Create a Qarray from raw data.

Handles shape normalisation, dimension inference, and tidying of small values.

State vectors are stored with their Hilbert space on a single trailing axis — a ket/bra of dimension N has data shape bdims + (N,) (never (N,1) / (1,N)). Operators keep the last two axes: bdims + (M, N). The ket/bra/oper distinction lives in _qdims, not in the data shape.

Legacy (N,1) / (1,N) inputs are still accepted and are squeezed to (N,) on the way in. Because a 1‑D (N,) array (or a square (N,N) batch of vectors) is shape-ambiguous, pass qtype to be explicit.

Parameters:

Name	Type	Description	Default
data		Input data array (dense array-like or sparse.BCOO).	required
dims		Quantum dimensions as ((row_dims...), (col_dims...)) or, as a convenience, a flat list of Hilbert space sizes. Inferred from data shape when None.	None
bdims		Tuple of batch dimension sizes. Inferred from the leading dimensions of data when None.	None
qtype		Optional quantum type — "ket" / "bra" / "oper" (or a :class:Qtypes member). When given it is authoritative and resolves shape ambiguity (e.g. a bare (N,) array with qtype="ket" behaves like a legacy (N,1) input).	None
implementation		Storage backend — QarrayImplType.DENSE (default) or QarrayImplType.SPARSE_BCOO, or the equivalent string "dense" / "sparse_bcoo".	DENSE

Returns:

Type	Description
	A new Qarray backed by the requested implementation.

Source code in jaxquantum/core/qarray.py

@classmethod
def create(cls, data, dims=None, bdims=None, qtype=None, implementation=QarrayImplType.DENSE):
    """Create a ``Qarray`` from raw data.

    Handles shape normalisation, dimension inference, and tidying of small
    values.

    State vectors are stored with their Hilbert space on a **single trailing
    axis** — a ket/bra of dimension ``N`` has data shape ``bdims + (N,)``
    (never ``(N,1)`` / ``(1,N)``).  Operators keep the last two axes:
    ``bdims + (M, N)``.  The ket/bra/oper distinction lives in ``_qdims``,
    not in the data shape.

    Legacy ``(N,1)`` / ``(1,N)`` inputs are still accepted and are squeezed
    to ``(N,)`` on the way in.  Because a 1‑D ``(N,)`` array (or a square
    ``(N,N)`` batch of vectors) is shape-ambiguous, pass ``qtype`` to be
    explicit.

    Args:
        data: Input data array (dense array-like or ``sparse.BCOO``).
        dims: Quantum dimensions as ``((row_dims...), (col_dims...))`` or, as
            a convenience, a flat list of Hilbert space sizes.  Inferred from
            *data* shape when ``None``.
        bdims: Tuple of batch dimension sizes.  Inferred from the leading
            dimensions of *data* when ``None``.
        qtype: Optional quantum type — ``"ket"`` / ``"bra"`` / ``"oper"`` (or
            a :class:`Qtypes` member).  When given it is authoritative and
            resolves shape ambiguity (e.g. a bare ``(N,)`` array with
            ``qtype="ket"`` behaves like a legacy ``(N,1)`` input).
        implementation: Storage backend — ``QarrayImplType.DENSE``
            (default) or ``QarrayImplType.SPARSE_BCOO``, or the equivalent
            string ``"dense"`` / ``"sparse_bcoo"``.

    Returns:
        A new ``Qarray`` backed by the requested implementation.
    """
    data = robust_asarray(data)

    # Whether the caller explicitly requested a qtype (vs. it being derived
    # from a full ``dims`` tuple). An explicit qtype is authoritative and
    # mismatches with the data raise a clear error.
    qtype_explicit = qtype is not None
    if qtype is not None and not isinstance(qtype, Qtypes):
        # Accepts "ket" / "bra" / "oper" (case-sensitive enum value).
        qtype = Qtypes.from_str(qtype)

    # Empty array (e.g. an empty list of operators) — keep legacy behaviour.
    if len(data.shape) == 1 and data.shape[0] == 0:
        dims = ((), ()) if dims is None else (tuple(dims[0]), tuple(dims[1]))
        bdims = (0,) if bdims is None else tuple(bdims)
        check_dims(dims, bdims, data.shape)
        qdims = Qdims(dims)
        impl_class = QarrayImplType(implementation).get_impl_class()
        impl = impl_class.from_data(data)
        impl = impl.tidy_up(SETTINGS["auto_tidyup_atol"])
        return cls(impl, qdims, bdims)

    # Resolve full quantum dims + batch dims.
    dims_is_full = dims is not None and isinstance(dims[0], (list, tuple))
    if qtype is None and dims_is_full:
        qtype = Qtypes.from_dims((tuple(dims[0]), tuple(dims[1])))

    if qtype is not None:
        # Explicit (or dims-derived) qtype: unambiguous layout.
        data, dims, bdims = cls._qtype_layout(data, dims, bdims, qtype)
    else:
        # Auto: legacy shape heuristics (1‑D → ket, square → oper, etc.).
        data, dims, bdims = cls._legacy_layout(data, dims, bdims)

    qdims = Qdims(dims)

    # Store vectors as bdims + (N,); operators as bdims + (M, N).
    if qdims.qtype == Qtypes.oper:
        space = (prod(dims[0]), prod(dims[1]))
    else:
        space = (prod(dims[0]) * prod(dims[1]),)

    target_shape = tuple(bdims) + space

    # An explicitly requested qtype is authoritative: error out if the data
    # cannot be interpreted as that type rather than silently coercing it.
    if qtype_explicit:
        if qdims.qtype != qtype:
            raise ValueError(
                f"Requested qtype={qtype} is incompatible with the resolved "
                f"dims {dims} (which is a {qdims.qtype})."
            )
        if prod(data.shape) != prod(target_shape):
            raise ValueError(
                f"Data of shape {tuple(data.shape)} is incompatible with "
                f"qtype={qtype} and dims={dims}: cannot reshape to "
                f"{target_shape}."
            )

    # Reshape to the canonical stored shape only when needed. Use varargs so
    # this works for dense arrays, BCOO, and SparseDiaData alike.
    if tuple(data.shape) != target_shape:
        data = data.reshape(*target_shape)

    check_dims(dims, bdims, data.shape)

    # NOTE: Constantly tidying up on Qarray creation might be a bit overkill.
    # It increases the compilation time, but only very slightly
    # increased the runtime of the jit compiled function.
    # We could instead use this tidy up where we think we need it.

    impl_class = QarrayImplType(implementation).get_impl_class()
    impl = impl_class.from_data(data)
    impl = impl.tidy_up(SETTINGS["auto_tidyup_atol"])

    return cls(impl, qdims, bdims)

dag()

Conjugate transpose of this array.

Source code in jaxquantum/core/qarray.py

def dag(self):
    """Conjugate transpose of this array."""
    return dag(self)

eigenenergies()

Eigenvalues of this operator.

Source code in jaxquantum/core/qarray.py

def eigenenergies(self):
    """Eigenvalues of this operator."""
    return eigenenergies(self)

eigenstates()

Eigenvalues and eigenstates of this operator.

Source code in jaxquantum/core/qarray.py

def eigenstates(self):
    """Eigenvalues and eigenstates of this operator."""
    return eigenstates(self)

eigenvalues()

Eigenvalues of this operator (alias for :meth:eigenenergies).

Source code in jaxquantum/core/qarray.py

def eigenvalues(self):
    """Eigenvalues of this operator (alias for :meth:`eigenenergies`)."""
    return eigenenergies(self)

expm()

Matrix exponential.

Source code in jaxquantum/core/qarray.py

def expm(self):
    """Matrix exponential."""
    return expm(self)

frobenius_norm()

Compute the Frobenius norm directly from the implementation.

Returns:

Type	Description
	The Frobenius norm as a scalar.

Source code in jaxquantum/core/qarray.py

def frobenius_norm(self):
    """Compute the Frobenius norm directly from the implementation.

    Returns:
        The Frobenius norm as a scalar.
    """
    return self._impl.frobenius_norm()

from_array(qarr_arr) classmethod

from_array(qarr_arr: 'Qarray[DenseImpl]') -> 'Qarray[DenseImpl]'

from_array(qarr_arr: 'Qarray[SparseBCOOImpl]') -> 'Qarray[SparseBCOOImpl]'

Create a Qarray from a (possibly nested) list of Qarray objects.

Parameters:

Name	Type	Description	Default
qarr_arr		A Qarray (returned as-is) or a nested list of Qarray objects.	required

Returns:

Type	Description
Qarray	A Qarray with batch dimensions matching the nesting structure
Qarray	of qarr_arr.

Source code in jaxquantum/core/qarray.py

@classmethod
def from_array(cls, qarr_arr) -> Qarray:
    """Create a ``Qarray`` from a (possibly nested) list of ``Qarray`` objects.

    Args:
        qarr_arr: A ``Qarray`` (returned as-is) or a nested list of
            ``Qarray`` objects.

    Returns:
        A ``Qarray`` with batch dimensions matching the nesting structure
        of *qarr_arr*.
    """
    if isinstance(qarr_arr, Qarray):
        return qarr_arr

    bdims = ()
    lvl = qarr_arr
    while not isinstance(lvl, Qarray):
        bdims = bdims + (len(lvl),)
        if len(lvl) > 0:
            lvl = lvl[0]
        else:
            break

    def flat(lis):
        flatList = []
        for element in lis:
            if type(element) is list:
                flatList += flat(element)
            else:
                flatList.append(element)
        return flatList

    qarr_list = flat(qarr_arr)
    qarr = cls.from_list(qarr_list)
    qarr = qarr.reshape_bdims(*bdims)
    return qarr

from_list(qarr_list, qtype=None) classmethod

from_list(qarr_list: List['Qarray[DenseImpl]'], qtype=None) -> 'Qarray[DenseImpl]'

from_list(qarr_list: List['Qarray[SparseBCOOImpl]'], qtype=None) -> 'Qarray[SparseBCOOImpl]'

Create a batched Qarray from a list of same-shaped Qarray objects.

The output implementation is determined by the element with the highest PROMOTION_ORDER: if all inputs are sparse the result is sparse; if any input is dense (or types are mixed) all inputs are promoted to dense and the result is dense.

Works for kets/bras (stacked into (len, *bdims, N)) as well as operators, regardless of whether the elements were originally created from (N,) or legacy (N,1) / (1,N) arrays — they are all stored as (N,) vectors by the time they reach here.

Parameters:

Name	Type	Description	Default
qarr_list	List[Qarray]	List of Qarray objects with identical dims and bdims. May be empty.	required
qtype		Optional quantum type ("ket"/"bra"/"oper" or a Qtypes) forwarded to :meth:create. Defaults to the qtype of the first element, which is the correct choice in all normal cases; pass it explicitly to override or to be defensive.	None

Returns:

Type	Description
Qarray	A Qarray with an extra leading batch dimension of size
Qarray	len(qarr_list).

Raises:

Type	Description
ValueError	If the elements have mismatched dims or bdims.

Source code in jaxquantum/core/qarray.py

@classmethod
def from_list(cls, qarr_list: List[Qarray], qtype=None) -> Qarray:
    """Create a batched ``Qarray`` from a list of same-shaped ``Qarray`` objects.

    The output implementation is determined by the element with the highest
    ``PROMOTION_ORDER``: if all inputs are sparse the result is sparse; if
    any input is dense (or types are mixed) all inputs are promoted to dense
    and the result is dense.

    Works for kets/bras (stacked into ``(len, *bdims, N)``) as well as
    operators, regardless of whether the elements were originally created
    from ``(N,)`` or legacy ``(N,1)`` / ``(1,N)`` arrays — they are all
    stored as ``(N,)`` vectors by the time they reach here.

    Args:
        qarr_list: List of ``Qarray`` objects with identical ``dims`` and
            ``bdims``.  May be empty.
        qtype: Optional quantum type ("ket"/"bra"/"oper" or a ``Qtypes``)
            forwarded to :meth:`create`.  Defaults to the qtype of the
            first element, which is the correct choice in all normal cases;
            pass it explicitly to override or to be defensive.

    Returns:
        A ``Qarray`` with an extra leading batch dimension of size
        ``len(qarr_list)``.

    Raises:
        ValueError: If the elements have mismatched ``dims`` or ``bdims``.
    """
    if len(qarr_list) == 0:
        dims = ((), ())
        bdims = (0,)
        return cls.create(jnp.array([]), dims=dims, bdims=bdims)

    dims = qarr_list[0].dims
    bdims = qarr_list[0].bdims

    if not all(qarr.dims == dims and qarr.bdims == bdims for qarr in qarr_list):
        raise ValueError("All Qarrays in the list must have the same dimensions.")

    if qtype is None:
        qtype = qarr_list[0].qtype

    new_bdims = (len(qarr_list),) + bdims

    # Pick the target type: highest PROMOTION_ORDER wins (dense beats sparse).
    target_impl_type = max(
        (q.impl_type for q in qarr_list),
        key=lambda t: t.get_impl_class().PROMOTION_ORDER,
    )

    if target_impl_type == QarrayImplType.SPARSE_DIA:
        # All inputs are SparseDIA — batch without densifying.
        # Compute union of offsets across all operators, then remap each
        # operator's _diags rows into the union shape and stack.
        from jaxquantum.core.sparse_dia import SparseDiaData  # lazy to avoid circular
        union_offsets = tuple(sorted(
            set().union(*[set(q._impl._offsets) for q in qarr_list])
        ))
        union_idx = {k: i for i, k in enumerate(union_offsets)}
        n = qarr_list[0]._impl._diags.shape[-1]
        dtype = jnp.result_type(*[q._impl._diags.dtype for q in qarr_list])
        remapped = []
        for q in qarr_list:
            row = jnp.zeros((len(union_offsets), n), dtype=dtype)
            for i_src, k in enumerate(q._impl._offsets):
                row = row.at[union_idx[k], :].set(q._impl._diags[i_src, :])
            remapped.append(row)
        stacked = jnp.stack(remapped, axis=0)  # (n_ops, n_union_diags, N)
        raw = SparseDiaData(offsets=union_offsets, diags=stacked)
        return cls.create(raw, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.SPARSE_DIA)

    if target_impl_type == QarrayImplType.SPARSE_BCOO:
        # All inputs are sparse BCOO — stack via dense intermediates then re-sparsify.
        data = jnp.array([q.data.todense() for q in qarr_list])
        return cls.create(data, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.SPARSE_BCOO)

    # Target is dense: promote any sparse inputs before stacking.
    data = jnp.array([q.to_dense().data for q in qarr_list])
    return cls.create(data, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.DENSE)

from_sparse_bcoo(data, dims=None, bdims=None) classmethod

from_sparse_bcoo(data, dims=None, bdims=None) -> 'Qarray[SparseBCOOImpl]'

Create a Qarray directly from a sparse BCOO array without densifying.

Parameters:

Name	Type	Description	Default
data		A sparse.BCOO or array-like to store as sparse BCOO.	required
dims		Quantum dimensions. Inferred when None.	None
bdims		Batch dimensions. Inferred when None.	None

Returns:

Type	Description
	A Qarray[SparseBCOOImpl].

Source code in jaxquantum/core/qarray.py

@classmethod
def from_sparse_bcoo(cls, data, dims=None, bdims=None):
    """Create a ``Qarray`` directly from a sparse BCOO array without densifying.

    Args:
        data: A ``sparse.BCOO`` or array-like to store as sparse BCOO.
        dims: Quantum dimensions.  Inferred when ``None``.
        bdims: Batch dimensions.  Inferred when ``None``.

    Returns:
        A ``Qarray[SparseBCOOImpl]``.
    """
    return cls.create(data, dims=dims, bdims=bdims, implementation=QarrayImplType.SPARSE_BCOO)

from_sparse_dia(data, dims=None, bdims=None) classmethod

Create a SparseDIA-backed Qarray.

Accepts either a dense array-like (diagonals are auto-detected) or a :class:~jaxquantum.core.sparse_dia.SparseDiaData container.

Parameters:

Name	Type	Description	Default
data		Dense array of shape (*batch, n, n) or a SparseDiaData.	required
dims		Quantum dimensions ((row_dims,), (col_dims,)).	None
bdims		Batch dimension sizes.	None

Returns:

Type	Description
'Qarray'	A Qarray backed by SparseDiaImpl.

Source code in jaxquantum/core/qarray.py

@classmethod
def from_sparse_dia(cls, data, dims=None, bdims=None) -> "Qarray":
    """Create a SparseDIA-backed ``Qarray``.

    Accepts either a dense array-like (diagonals are auto-detected) or a
    :class:`~jaxquantum.core.sparse_dia.SparseDiaData` container.

    Args:
        data: Dense array of shape (*batch, n, n) or a ``SparseDiaData``.
        dims: Quantum dimensions ``((row_dims,), (col_dims,))``.
        bdims: Batch dimension sizes.

    Returns:
        A ``Qarray`` backed by ``SparseDiaImpl``.
    """
    return cls.create(data, dims=dims, bdims=bdims, implementation=QarrayImplType.SPARSE_DIA)

imag()

Element-wise imaginary part.

Returns:

Type	Description
	A new Qarray containing the imaginary parts of each element.

Source code in jaxquantum/core/qarray.py

def imag(self):
    """Element-wise imaginary part.

    Returns:
        A new ``Qarray`` containing the imaginary parts of each element.
    """
    new_impl = self._impl.imag()
    return Qarray._from_impl(new_impl, self._qdims)

is_dm()

Return True if this array is an operator (density-matrix type).

Source code in jaxquantum/core/qarray.py

def is_dm(self):
    """Return True if this array is an operator (density-matrix type)."""
    return self.qtype == Qtypes.oper

is_vec()

Return True if this array is a ket or bra.

Source code in jaxquantum/core/qarray.py

def is_vec(self):
    """Return True if this array is a ket or bra."""
    return self.qtype == Qtypes.ket or self.qtype == Qtypes.bra

keep_only_diag_elements()

Zero out all off-diagonal elements.

Source code in jaxquantum/core/qarray.py

def keep_only_diag_elements(self):
    """Zero out all off-diagonal elements."""
    return keep_only_diag_elements(self)

norm()

Compute the norm of this array.

Source code in jaxquantum/core/qarray.py

def norm(self):
    """Compute the norm of this array."""
    return norm(self)

powm(n)

Matrix power.

Parameters:

Name	Type	Description	Default
n		Exponent (integer or float).	required

Returns:

Type	Description
	This array raised to the n-th matrix power.

Source code in jaxquantum/core/qarray.py

def powm(self, n):
    """Matrix power.

    Args:
        n: Exponent (integer or float).

    Returns:
        This array raised to the *n*-th matrix power.
    """
    return powm(self, n)

ptrace(indx)

Partial trace over subsystem indx.

Parameters:

Name	Type	Description	Default
indx		Index of the subsystem to trace out.	required

Returns:

Type	Description
	Reduced density matrix.

Source code in jaxquantum/core/qarray.py

def ptrace(self, indx):
    """Partial trace over subsystem *indx*.

    Args:
        indx: Index of the subsystem to trace out.

    Returns:
        Reduced density matrix.
    """
    return ptrace(self, indx)

real()

Element-wise real part.

Returns:

Type	Description
	A new Qarray containing the real parts of each element.

Source code in jaxquantum/core/qarray.py

def real(self):
    """Element-wise real part.

    Returns:
        A new ``Qarray`` containing the real parts of each element.
    """
    new_impl = self._impl.real()
    return Qarray._from_impl(new_impl, self._qdims)

reshape_bdims(*args)

Reshape the batch dimensions of this Qarray.

Parameters:

Name	Type	Description	Default
*args		New batch dimension sizes.	()

Returns:

Type	Description
	A new Qarray with the requested batch shape.

Source code in jaxquantum/core/qarray.py

def reshape_bdims(self, *args):
    """Reshape the batch dimensions of this ``Qarray``.

    Args:
        *args: New batch dimension sizes.

    Returns:
        A new ``Qarray`` with the requested batch shape.
    """
    new_bdims = tuple(args)

    if prod(new_bdims) == 0:
        new_shape = new_bdims
    elif self.qtype == Qtypes.oper:
        new_shape = new_bdims + (prod(self.dims[0]), -1)
    else:
        # Vectors keep a single trailing space axis (no (N,1)).
        new_shape = new_bdims + (prod(self.dims[0]) * prod(self.dims[1]),)

    impl = type(self._impl).from_data(self.data.reshape(new_shape))
    return Qarray._from_impl(
        impl,
        self._qdims,
        new_bdims,
    )

reshape_qdims(*args)

Reshape the quantum dimensions of the Qarray.

Note that this does not take in qdims but rather the new Hilbert space dimensions.

Parameters:

Name	Type	Description	Default
*args		New Hilbert dimensions for the Qarray.	()

Returns:

Name	Type	Description
Qarray		reshaped Qarray.

Source code in jaxquantum/core/qarray.py

def reshape_qdims(self, *args):
    """Reshape the quantum dimensions of the Qarray.

    Note that this does not take in qdims but rather the new Hilbert space
    dimensions.

    Args:
        *args: New Hilbert dimensions for the Qarray.

    Returns:
        Qarray: reshaped Qarray.
    """

    new_space_dims = tuple(args)
    current_space_dims = self.space_dims
    assert prod(new_space_dims) == prod(current_space_dims)

    new_qdims = self.space_to_qdims(new_space_dims)
    return Qarray._from_impl(self._impl, Qdims(new_qdims), self._bdims)

resize(new_shape)

Resize the Qarray to a new shape.

TODO: review and maybe deprecate this method.

Parameters:

Name	Type	Description	Default
new_shape		Target shape tuple.	required

Returns:

Type	Description
	A new Qarray with data resized via jnp.resize.

Source code in jaxquantum/core/qarray.py

def resize(self, new_shape):
    """Resize the Qarray to a new shape.

    TODO: review and maybe deprecate this method.

    Args:
        new_shape: Target shape tuple.

    Returns:
        A new ``Qarray`` with data resized via ``jnp.resize``.
    """
    dims = self.dims
    data = jnp.resize(self.data, new_shape)
    # Preserve implementation type
    implementation = self.impl_type
    return Qarray.create(
        data,
        dims=dims,
        implementation=implementation,
    )

sinm()

Matrix sine.

Source code in jaxquantum/core/qarray.py

def sinm(self):
    """Matrix sine."""
    return sinm(self)

space_to_qdims(space_dims)

Convert Hilbert space dimensions to full quantum dims tuple.

Parameters:

Name	Type	Description	Default
space_dims	List[int]	Sequence of per-subsystem Hilbert space sizes, or a full ((row_dims), (col_dims)) tuple (returned unchanged).	required

Returns:

Type	Description
	A ((row_dims...), (col_dims...)) tuple.

Raises:

Type	Description
ValueError	If self.qtype is not ket, bra, or oper.

Source code in jaxquantum/core/qarray.py

def space_to_qdims(self, space_dims: List[int]):
    """Convert Hilbert space dimensions to full quantum dims tuple.

    Args:
        space_dims: Sequence of per-subsystem Hilbert space sizes, or a
            full ``((row_dims), (col_dims))`` tuple (returned unchanged).

    Returns:
        A ``((row_dims...), (col_dims...))`` tuple.

    Raises:
        ValueError: If ``self.qtype`` is not ket, bra, or oper.
    """
    if isinstance(space_dims[0], (list, tuple)):
        return space_dims

    if self.qtype == Qtypes.oper:
        return (tuple(space_dims), tuple(space_dims))
    elif self.qtype == Qtypes.ket:
        return (tuple(space_dims), tuple([1 for _ in range(len(space_dims))]))
    elif self.qtype == Qtypes.bra:
        return (tuple([1 for _ in range(len(space_dims))]), tuple(space_dims))
    else:
        raise ValueError("Unsupported qtype for space_to_qdims conversion.")

to_dense()

Return a dense-backed copy of this array.

If the array is already dense, returns self unchanged.

Returns:

Type	Description
'Qarray[DenseImpl]'	A Qarray[DenseImpl].

Source code in jaxquantum/core/qarray.py

def to_dense(self) -> "Qarray[DenseImpl]":
    """Return a dense-backed copy of this array.

    If the array is already dense, returns self unchanged.

    Returns:
        A ``Qarray[DenseImpl]``.
    """
    if self.is_dense:
        return self
    new_impl = self._impl.to_dense()
    return Qarray(new_impl, self._qdims, self._bdims)

to_dm()

Convert a ket to a density matrix via outer product.

Source code in jaxquantum/core/qarray.py

def to_dm(self):
    """Convert a ket to a density matrix via outer product."""
    return ket2dm(self)

to_ket()

Convert a bra to a ket (no-op for kets).

Source code in jaxquantum/core/qarray.py

def to_ket(self):
    """Convert a bra to a ket (no-op for kets)."""
    return to_ket(self)

to_sparse_bcoo()

Return a BCOO-sparse-backed copy of this array.

If the array is already sparse BCOO, returns self unchanged.

Returns:

Type	Description
'Qarray[SparseBCOOImpl]'	A Qarray[SparseBCOOImpl].

Source code in jaxquantum/core/qarray.py

def to_sparse_bcoo(self) -> "Qarray[SparseBCOOImpl]":
    """Return a BCOO-sparse-backed copy of this array.

    If the array is already sparse BCOO, returns self unchanged.

    Returns:
        A ``Qarray[SparseBCOOImpl]``.
    """
    if self.is_sparse_bcoo:
        return self
    new_impl = self._impl.to_sparse_bcoo()
    return Qarray(new_impl, self._qdims, self._bdims)

to_sparse_dia()

Return a SparseDIA-backed copy of this array.

If the array is already SparseDIA, returns self unchanged.

Returns:

Type	Description
'Qarray'	A Qarray[SparseDiaImpl].

Source code in jaxquantum/core/qarray.py

def to_sparse_dia(self) -> "Qarray":
    """Return a SparseDIA-backed copy of this array.

    If the array is already SparseDIA, returns self unchanged.

    Returns:
        A ``Qarray[SparseDiaImpl]``.
    """
    if self.is_sparse_dia:
        return self
    new_impl = self._impl.to_sparse_dia()
    return Qarray(new_impl, self._qdims, self._bdims)

tr(**kwargs)

Full trace.

Source code in jaxquantum/core/qarray.py

def tr(self, **kwargs):
    """Full trace."""
    return tr(self, **kwargs)

trace(**kwargs)

Full trace (alias for :meth:tr).

Source code in jaxquantum/core/qarray.py

def trace(self, **kwargs):
    """Full trace (alias for :meth:`tr`)."""
    return tr(self, **kwargs)

transpose(*args)

Transpose subsystem indices.

Source code in jaxquantum/core/qarray.py

def transpose(self, *args):
    """Transpose subsystem indices."""
    return transpose(self, *args)

unit()

Return the normalised (unit-norm) version of this array.

Source code in jaxquantum/core/qarray.py

def unit(self):
    """Return the normalised (unit-norm) version of this array."""
    return unit(self)

QarrayImpl

Bases: ABC

Abstract base class defining the interface every storage backend must implement.

A QarrayImpl wraps a raw data array (dense jnp.ndarray or sparse BCOO) and provides the mathematical primitives used by Qarray. Concrete subclasses must implement every @abstractmethod.

Attributes:

Name	Type	Description
PROMOTION_ORDER	int	Integer priority used by _coerce to decide which side to promote when operands have different types. Higher means "more general" (DenseImpl = 1, SparseBCOOImpl = 0).

Source code in jaxquantum/core/qarray.py

class QarrayImpl(ABC):
    """Abstract base class defining the interface every storage backend must implement.

    A ``QarrayImpl`` wraps a raw data array (dense ``jnp.ndarray`` or sparse
    ``BCOO``) and provides the mathematical primitives used by ``Qarray``.
    Concrete subclasses must implement every ``@abstractmethod``.

    Attributes:
        PROMOTION_ORDER: Integer priority used by ``_coerce`` to decide which
            side to promote when operands have different types.  Higher means
            "more general" (``DenseImpl = 1``, ``SparseBCOOImpl = 0``).
    """

    PROMOTION_ORDER: int = 0  # override in subclasses; higher = more general
    # Current hierarchy: SparseDiaImpl=0, SparseBCOOImpl=1, DenseImpl=2

    @abstractmethod
    def get_data(self) -> Array:
        """Return the underlying raw data array."""
        pass

    @property
    def data(self) -> Array:
        """The underlying raw data array."""
        return self.get_data()

    @property
    def impl_type(self) -> QarrayImplType:
        """The ``QarrayImplType`` member corresponding to this instance."""
        return QarrayImplType.from_impl_class(type(self))

    @classmethod
    @abstractmethod
    def from_data(cls, data) -> "QarrayImpl":
        """Wrap raw data in this impl type.

        Args:
            data: Raw array data (dense ``jnp.ndarray`` or ``sparse.BCOO``).

        Returns:
            A new instance of this implementation wrapping *data*.
        """
        pass

    @abstractmethod
    def matmul(self, other: "QarrayImpl") -> "QarrayImpl":
        """Matrix multiplication with *other*.

        Args:
            other: Right-hand operand.

        Returns:
            Result of ``self @ other`` as a ``QarrayImpl``.
        """
        pass

    @abstractmethod
    def add(self, other: "QarrayImpl") -> "QarrayImpl":
        """Element-wise addition with *other*.

        Args:
            other: Right-hand operand.

        Returns:
            Result of ``self + other`` as a ``QarrayImpl``.
        """
        pass

    @abstractmethod
    def sub(self, other: "QarrayImpl") -> "QarrayImpl":
        """Element-wise subtraction of *other*.

        Args:
            other: Right-hand operand.

        Returns:
            Result of ``self - other`` as a ``QarrayImpl``.
        """
        pass

    @abstractmethod
    def mul(self, scalar) -> "QarrayImpl":
        """Scalar multiplication.

        Args:
            scalar: Scalar value to multiply by.

        Returns:
            Result of ``scalar * self`` as a ``QarrayImpl``.
        """
        pass

    @abstractmethod
    def dag(self) -> "QarrayImpl":
        """Conjugate transpose.

        Returns:
            The conjugate transpose of this array as a ``QarrayImpl``.
        """
        pass

    @abstractmethod
    def to_dense(self) -> "DenseImpl":
        """Convert to a ``DenseImpl``.

        Returns:
            A ``DenseImpl`` wrapping the same data.
        """
        pass

    @abstractmethod
    def to_sparse_bcoo(self) -> "SparseBCOOImpl":
        """Convert to a ``SparseBCOOImpl`` (BCOO).

        Returns:
            A ``SparseBCOOImpl`` wrapping the same data.
        """
        pass

    def to_sparse_dia(self) -> "QarrayImpl":
        """Convert to a ``SparseDiaImpl``.

        Default implementation goes through dense and auto-detects diagonals.
        Subclasses may override for a more direct path.

        Returns:
            A ``SparseDiaImpl`` wrapping the same data.
        """
        # Import here to avoid circular imports at module load time
        from jaxquantum.core.sparse_dia import SparseDiaImpl
        return SparseDiaImpl.from_data(self.to_dense()._data)

    @abstractmethod
    def shape(self) -> tuple:
        """Shape of the underlying data array.

        Returns:
            Tuple of dimension sizes.
        """
        pass

    @abstractmethod
    def dtype(self):
        """Data type of the underlying array.

        Returns:
            A numpy/JAX dtype object.
        """
        pass

    @abstractmethod
    def __deepcopy__(self, memo=None):
        pass

    @abstractmethod
    def tidy_up(self, atol):
        """Zero out values whose magnitude is below *atol*.

        Args:
            atol: Absolute tolerance threshold.

        Returns:
            A new ``QarrayImpl`` with small values zeroed.
        """
        pass

    @abstractmethod
    def kron(self, other: "QarrayImpl") -> "QarrayImpl":
        """Kronecker (tensor) product with another implementation.

        Args:
            other: Right-hand operand.  Mixed-type pairs are handled by
                ``_coerce`` — the result has the higher ``PROMOTION_ORDER``
                type (dense wins over sparse).

        Returns:
            A new ``QarrayImpl`` containing the Kronecker product.
        """
        pass

    @classmethod
    @abstractmethod
    def _eye_data(cls, n: int, dtype=None):
        """Create identity matrix data of size n.

        Args:
            n: Matrix size.
            dtype: Optional data type for the identity entries.

        Returns:
            Raw identity matrix data in the format appropriate for this impl.
        """
        pass

    @classmethod
    def _scaled_identity(cls, n: int, scalar, dtype=None) -> QarrayImpl:
        """Create an identity scaled over any leading scalar batch axes."""
        scalar = jnp.asarray(scalar) + 0.0j
        if scalar.ndim:
            scalar = scalar.reshape(scalar.shape + (1, 1))
        return cls.from_data(cls._eye_data(n, dtype=dtype) * scalar)

    @classmethod
    @abstractmethod
    def can_handle_data(cls, arr) -> bool:
        """Return True if *arr* is a raw data type natively handled by this impl.

        Used by the module-level :func:`dag_data` dispatcher to route raw
        arrays to the correct backend without any isinstance chain outside the
        impl classes.

        Args:
            arr: Raw array — e.g. ``jnp.ndarray`` for ``DenseImpl`` or
                ``sparse.BCOO`` for ``SparseBCOOImpl``.

        Returns:
            True if this impl can operate on *arr* without conversion.
        """
        pass

    @classmethod
    @abstractmethod
    def dag_data(cls, arr):
        """Conjugate transpose of raw data in this impl's native format.

        Implementations must handle batched arrays (last two axes are
        swapped) and must not densify sparse arrays.

        Args:
            arr: Raw array in this impl's native format.

        Returns:
            Conjugate transpose with the last two axes swapped.
        """
        pass

    def _promote_to(self, target_cls: type) -> "QarrayImpl":
        """Convert this impl to *target_cls* by passing through dense.

        Args:
            target_cls: The destination ``QarrayImpl`` subclass.

        Returns:
            An instance of *target_cls* holding equivalent data.
        """
        if isinstance(self, target_cls):
            return self
        return target_cls.from_data(self.to_dense()._data)

    def _coerce(self, other: "QarrayImpl") -> "tuple[QarrayImpl, QarrayImpl]":
        """Coerce *self* and *other* to the same implementation type.

        The impl type with the higher ``PROMOTION_ORDER`` wins; the other side
        is promoted via :meth:`_promote_to`.

        Args:
            other: The other operand.

        Returns:
            A pair ``(a, b)`` of the same ``QarrayImpl`` subclass, suitable
            for a binary operation.
        """
        if type(self) is type(other):
            return self, other
        if self.PROMOTION_ORDER >= other.PROMOTION_ORDER:
            return self, other._promote_to(type(self))
        return self._promote_to(type(other)), other

data property

The underlying raw data array.

impl_type property

The QarrayImplType member corresponding to this instance.

add(other) abstractmethod

Element-wise addition with other.

Parameters:

Name	Type	Description	Default
other	'QarrayImpl'	Right-hand operand.	required

Returns:

Type	Description
'QarrayImpl'	Result of self + other as a QarrayImpl.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def add(self, other: "QarrayImpl") -> "QarrayImpl":
    """Element-wise addition with *other*.

    Args:
        other: Right-hand operand.

    Returns:
        Result of ``self + other`` as a ``QarrayImpl``.
    """
    pass

can_handle_data(arr) abstractmethod classmethod

Return True if arr is a raw data type natively handled by this impl.

Used by the module-level :func:dag_data dispatcher to route raw arrays to the correct backend without any isinstance chain outside the impl classes.

Parameters:

Name	Type	Description	Default
arr		Raw array — e.g. jnp.ndarray for DenseImpl or sparse.BCOO for SparseBCOOImpl.	required

Returns:

Type	Description
bool	True if this impl can operate on arr without conversion.

Source code in jaxquantum/core/qarray.py

@classmethod
@abstractmethod
def can_handle_data(cls, arr) -> bool:
    """Return True if *arr* is a raw data type natively handled by this impl.

    Used by the module-level :func:`dag_data` dispatcher to route raw
    arrays to the correct backend without any isinstance chain outside the
    impl classes.

    Args:
        arr: Raw array — e.g. ``jnp.ndarray`` for ``DenseImpl`` or
            ``sparse.BCOO`` for ``SparseBCOOImpl``.

    Returns:
        True if this impl can operate on *arr* without conversion.
    """
    pass

dag() abstractmethod

Conjugate transpose.

Returns:

Type	Description
'QarrayImpl'	The conjugate transpose of this array as a QarrayImpl.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def dag(self) -> "QarrayImpl":
    """Conjugate transpose.

    Returns:
        The conjugate transpose of this array as a ``QarrayImpl``.
    """
    pass

dag_data(arr) abstractmethod classmethod

Conjugate transpose of raw data in this impl's native format.

Implementations must handle batched arrays (last two axes are swapped) and must not densify sparse arrays.

Parameters:

Name	Type	Description	Default
arr		Raw array in this impl's native format.	required

Returns:

Type	Description
	Conjugate transpose with the last two axes swapped.

Source code in jaxquantum/core/qarray.py

@classmethod
@abstractmethod
def dag_data(cls, arr):
    """Conjugate transpose of raw data in this impl's native format.

    Implementations must handle batched arrays (last two axes are
    swapped) and must not densify sparse arrays.

    Args:
        arr: Raw array in this impl's native format.

    Returns:
        Conjugate transpose with the last two axes swapped.
    """
    pass

dtype() abstractmethod

Data type of the underlying array.

Returns:

Type	Description
	A numpy/JAX dtype object.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def dtype(self):
    """Data type of the underlying array.

    Returns:
        A numpy/JAX dtype object.
    """
    pass

from_data(data) abstractmethod classmethod

Wrap raw data in this impl type.

Parameters:

Name	Type	Description	Default
data		Raw array data (dense jnp.ndarray or sparse.BCOO).	required

Returns:

Type	Description
'QarrayImpl'	A new instance of this implementation wrapping data.

Source code in jaxquantum/core/qarray.py

@classmethod
@abstractmethod
def from_data(cls, data) -> "QarrayImpl":
    """Wrap raw data in this impl type.

    Args:
        data: Raw array data (dense ``jnp.ndarray`` or ``sparse.BCOO``).

    Returns:
        A new instance of this implementation wrapping *data*.
    """
    pass

get_data() abstractmethod

Return the underlying raw data array.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def get_data(self) -> Array:
    """Return the underlying raw data array."""
    pass

kron(other) abstractmethod

Kronecker (tensor) product with another implementation.

Parameters:

Name	Type	Description	Default
other	'QarrayImpl'	Right-hand operand. Mixed-type pairs are handled by _coerce — the result has the higher PROMOTION_ORDER type (dense wins over sparse).	required

Returns:

Type	Description
'QarrayImpl'	A new QarrayImpl containing the Kronecker product.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def kron(self, other: "QarrayImpl") -> "QarrayImpl":
    """Kronecker (tensor) product with another implementation.

    Args:
        other: Right-hand operand.  Mixed-type pairs are handled by
            ``_coerce`` — the result has the higher ``PROMOTION_ORDER``
            type (dense wins over sparse).

    Returns:
        A new ``QarrayImpl`` containing the Kronecker product.
    """
    pass

matmul(other) abstractmethod

Matrix multiplication with other.

Parameters:

Name	Type	Description	Default
other	'QarrayImpl'	Right-hand operand.	required

Returns:

Type	Description
'QarrayImpl'	Result of self @ other as a QarrayImpl.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def matmul(self, other: "QarrayImpl") -> "QarrayImpl":
    """Matrix multiplication with *other*.

    Args:
        other: Right-hand operand.

    Returns:
        Result of ``self @ other`` as a ``QarrayImpl``.
    """
    pass

mul(scalar) abstractmethod

Scalar multiplication.

Parameters:

Name	Type	Description	Default
scalar		Scalar value to multiply by.	required

Returns:

Type	Description
'QarrayImpl'	Result of scalar * self as a QarrayImpl.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def mul(self, scalar) -> "QarrayImpl":
    """Scalar multiplication.

    Args:
        scalar: Scalar value to multiply by.

    Returns:
        Result of ``scalar * self`` as a ``QarrayImpl``.
    """
    pass

shape() abstractmethod

Shape of the underlying data array.

Returns:

Type	Description
tuple	Tuple of dimension sizes.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def shape(self) -> tuple:
    """Shape of the underlying data array.

    Returns:
        Tuple of dimension sizes.
    """
    pass

sub(other) abstractmethod

Element-wise subtraction of other.

Parameters:

Name	Type	Description	Default
other	'QarrayImpl'	Right-hand operand.	required

Returns:

Type	Description
'QarrayImpl'	Result of self - other as a QarrayImpl.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def sub(self, other: "QarrayImpl") -> "QarrayImpl":
    """Element-wise subtraction of *other*.

    Args:
        other: Right-hand operand.

    Returns:
        Result of ``self - other`` as a ``QarrayImpl``.
    """
    pass

tidy_up(atol) abstractmethod

Zero out values whose magnitude is below atol.

Parameters:

Name	Type	Description	Default
atol		Absolute tolerance threshold.	required

Returns:

Type	Description
	A new QarrayImpl with small values zeroed.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def tidy_up(self, atol):
    """Zero out values whose magnitude is below *atol*.

    Args:
        atol: Absolute tolerance threshold.

    Returns:
        A new ``QarrayImpl`` with small values zeroed.
    """
    pass

to_dense() abstractmethod

Convert to a DenseImpl.

Returns:

Type	Description
'DenseImpl'	A DenseImpl wrapping the same data.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def to_dense(self) -> "DenseImpl":
    """Convert to a ``DenseImpl``.

    Returns:
        A ``DenseImpl`` wrapping the same data.
    """
    pass

to_sparse_bcoo() abstractmethod

Convert to a SparseBCOOImpl (BCOO).

Returns:

Type	Description
'SparseBCOOImpl'	A SparseBCOOImpl wrapping the same data.

Source code in jaxquantum/core/qarray.py

@abstractmethod
def to_sparse_bcoo(self) -> "SparseBCOOImpl":
    """Convert to a ``SparseBCOOImpl`` (BCOO).

    Returns:
        A ``SparseBCOOImpl`` wrapping the same data.
    """
    pass

to_sparse_dia()

Convert to a SparseDiaImpl.

Default implementation goes through dense and auto-detects diagonals. Subclasses may override for a more direct path.

Returns:

Type	Description
'QarrayImpl'	A SparseDiaImpl wrapping the same data.

Source code in jaxquantum/core/qarray.py

def to_sparse_dia(self) -> "QarrayImpl":
    """Convert to a ``SparseDiaImpl``.

    Default implementation goes through dense and auto-detects diagonals.
    Subclasses may override for a more direct path.

    Returns:
        A ``SparseDiaImpl`` wrapping the same data.
    """
    # Import here to avoid circular imports at module load time
    from jaxquantum.core.sparse_dia import SparseDiaImpl
    return SparseDiaImpl.from_data(self.to_dense()._data)

QarrayImplType

Bases: Enum

Enumeration of available Qarray storage backends.

Each member maps one-to-one with a concrete QarrayImpl subclass. New backends should call QarrayImplType.register(MyImpl, QarrayImplType.MY_TYPE) immediately after defining their impl class.

Members

DENSE: Standard JAX dense array (jnp.ndarray). SPARSE_BCOO: JAX experimental BCOO sparse array. SPARSE_DIA: Diagonal sparse array.

Source code in jaxquantum/core/qarray.py

class QarrayImplType(Enum):
    """Enumeration of available Qarray storage backends.

    Each member maps one-to-one with a concrete ``QarrayImpl`` subclass.
    New backends should call ``QarrayImplType.register(MyImpl, QarrayImplType.MY_TYPE)``
    immediately after defining their impl class.

    Members:
        DENSE: Standard JAX dense array (``jnp.ndarray``).
        SPARSE_BCOO: JAX experimental BCOO sparse array.
        SPARSE_DIA: Diagonal sparse array.
    """

    DENSE = "dense"
    SPARSE_BCOO = "sparse_bcoo"
    SPARSE_DIA = "sparse_dia"

    @classmethod
    def register(cls, impl_class, member):
        """Register an implementation class with a QarrayImplType member.

        Args:
            impl_class: The concrete ``QarrayImpl`` subclass to register.
            member: The ``QarrayImplType`` enum member to associate with it.
        """
        _IMPL_REGISTRY[impl_class] = member

    @classmethod
    def has(cls, x) -> bool:
        """Return True if x corresponds to a member of QarrayImplType.

        Accepts an existing ``QarrayImplType`` member, a string equal to the
        member name or value (case-insensitive), or an implementation class
        (e.g. ``DenseImpl``, ``SparseBCOOImpl``) that has been registered.

        Args:
            x: Value to test — a ``QarrayImplType``, ``str``, or impl class.

        Returns:
            True if ``x`` maps to a known ``QarrayImplType`` member.
        """
        if isinstance(x, cls):
            return True

        if isinstance(x, str):
            xl = x.lower()
            return any(xl == member.value or xl == member.name.lower() for member in cls)

        # Try mapping from an implementation class to an enum member
        try:
            cls.from_impl_class(x)
            return True
        except Exception:  # noqa: BLE001
            return False

    @classmethod
    def from_impl_class(cls, impl_class) -> "QarrayImplType":
        """Return the ``QarrayImplType`` member associated with *impl_class*.

        Args:
            impl_class: A concrete ``QarrayImpl`` subclass that has been
                registered via :meth:`register`.

        Returns:
            The corresponding ``QarrayImplType`` member.

        Raises:
            ValueError: If *impl_class* is not in the registry.
        """
        if impl_class in _IMPL_REGISTRY:
            return _IMPL_REGISTRY[impl_class]
        raise ValueError(f"Unknown implementation class: {impl_class}")

    def get_impl_class(self):
        """Return the implementation class registered for this member.

        Returns:
            The concrete ``QarrayImpl`` subclass associated with this member.

        Raises:
            ValueError: If no class has been registered for this member.
        """
        for cls_key, member in _IMPL_REGISTRY.items():
            if member is self:
                return cls_key
        raise ValueError(f"No impl class registered for {self}")

from_impl_class(impl_class) classmethod

Return the QarrayImplType member associated with impl_class.

Parameters:

Name	Type	Description	Default
impl_class		A concrete QarrayImpl subclass that has been registered via :meth:register.	required

Returns:

Type	Description
'QarrayImplType'	The corresponding QarrayImplType member.

Raises:

Type	Description
ValueError	If impl_class is not in the registry.

Source code in jaxquantum/core/qarray.py

@classmethod
def from_impl_class(cls, impl_class) -> "QarrayImplType":
    """Return the ``QarrayImplType`` member associated with *impl_class*.

    Args:
        impl_class: A concrete ``QarrayImpl`` subclass that has been
            registered via :meth:`register`.

    Returns:
        The corresponding ``QarrayImplType`` member.

    Raises:
        ValueError: If *impl_class* is not in the registry.
    """
    if impl_class in _IMPL_REGISTRY:
        return _IMPL_REGISTRY[impl_class]
    raise ValueError(f"Unknown implementation class: {impl_class}")

get_impl_class()

Return the implementation class registered for this member.

Returns:

Type	Description
	The concrete QarrayImpl subclass associated with this member.

Raises:

Type	Description
ValueError	If no class has been registered for this member.

Source code in jaxquantum/core/qarray.py

def get_impl_class(self):
    """Return the implementation class registered for this member.

    Returns:
        The concrete ``QarrayImpl`` subclass associated with this member.

    Raises:
        ValueError: If no class has been registered for this member.
    """
    for cls_key, member in _IMPL_REGISTRY.items():
        if member is self:
            return cls_key
    raise ValueError(f"No impl class registered for {self}")

has(x) classmethod

Return True if x corresponds to a member of QarrayImplType.

Accepts an existing QarrayImplType member, a string equal to the member name or value (case-insensitive), or an implementation class (e.g. DenseImpl, SparseBCOOImpl) that has been registered.

Parameters:

Name	Type	Description	Default
x		Value to test — a QarrayImplType, str, or impl class.	required

Returns:

Type	Description
bool	True if x maps to a known QarrayImplType member.

Source code in jaxquantum/core/qarray.py

@classmethod
def has(cls, x) -> bool:
    """Return True if x corresponds to a member of QarrayImplType.

    Accepts an existing ``QarrayImplType`` member, a string equal to the
    member name or value (case-insensitive), or an implementation class
    (e.g. ``DenseImpl``, ``SparseBCOOImpl``) that has been registered.

    Args:
        x: Value to test — a ``QarrayImplType``, ``str``, or impl class.

    Returns:
        True if ``x`` maps to a known ``QarrayImplType`` member.
    """
    if isinstance(x, cls):
        return True

    if isinstance(x, str):
        xl = x.lower()
        return any(xl == member.value or xl == member.name.lower() for member in cls)

    # Try mapping from an implementation class to an enum member
    try:
        cls.from_impl_class(x)
        return True
    except Exception:  # noqa: BLE001
        return False

register(impl_class, member) classmethod

Register an implementation class with a QarrayImplType member.

Parameters:

Name	Type	Description	Default
impl_class		The concrete QarrayImpl subclass to register.	required
member		The QarrayImplType enum member to associate with it.	required

Source code in jaxquantum/core/qarray.py

@classmethod
def register(cls, impl_class, member):
    """Register an implementation class with a QarrayImplType member.

    Args:
        impl_class: The concrete ``QarrayImpl`` subclass to register.
        member: The ``QarrayImplType`` enum member to associate with it.
    """
    _IMPL_REGISTRY[impl_class] = member

QuantumStateTomography

Source code in jaxquantum/core/measurements.py

class QuantumStateTomography:
    def __init__(
        self,
        rho_guess: Qarray,
        measurement_basis: Qarray,
        measurement_results: jnp.ndarray,
        complete_basis: Optional[Qarray] = None,
        true_rho: Optional[Qarray] = None,
    ):
        """
        Reconstruct a quantum state from measurement results using quantum state tomography.
        The tomography can be performed either by direct inversion or by maximum likelihood estimation.

        Args:
            rho_guess (Qarray): The initial guess for the quantum state.
            measurement_basis (Qarray): The basis in which measurements are performed.
            measurement_results (jnp.ndarray): The results of the measurements.
            complete_basis (Optional[Qarray]): The complete basis for state 
            reconstruction used when using direct inversion. 
            Defaults to the measurement basis if not provided.
            true_rho (Optional[Qarray]): The true quantum state, if known.

        """
        self.rho_guess = rho_guess.data
        self.measurement_basis = measurement_basis.data
        self.measurement_results = measurement_results
        self.complete_basis = (
            complete_basis.data
            if (complete_basis is not None)
            else measurement_basis.data
        )
        self.true_rho = true_rho
        self._result = None

    @property
    def result(self) -> Optional[MLETomographyResult]:
        return self._result


    def quantum_state_tomography_mle(
        self, L1_reg_strength: float = 0.0, epochs: int = 10000, lr: float = 5e-3
    ) -> MLETomographyResult:
        """Perform quantum state tomography using maximum likelihood 
        estimation (MLE).

        This method reconstructs the quantum state from measurement results 
        by optimizing
        a likelihood function using gradient descent. The optimization 
        ensures the 
        resulting density matrix is positive semi-definite with trace 1.

        Args:
            L1_reg_strength (float, optional): Strength of L1 
            regularization. Defaults to 0.0.
            epochs (int, optional): Number of optimization iterations. 
            Defaults to 10000.
            lr (float, optional): Learning rate for the Adam optimizer. 
            Defaults to 5e-3.

        Returns:
            MLETomographyResult: Named tuple containing:
                - rho: Reconstructed quantum state as Qarray
                - params_history: List of parameter values during optimization
                - loss_history: List of loss values during optimization
                - grads_history: List of gradient values during optimization
                - infidelity_history: List of infidelities if true_rho was 
                provided, None otherwise
        """

        dim = self.rho_guess.shape[0]
        optimizer = optax.adam(lr)

        # Initialize parameters from the initial guess for the density matrix
        params = _parametrize_density_matrix(self.rho_guess, dim)
        opt_state = optimizer.init(params)

        compute_infidelity_flag = self.true_rho is not None

        # Provide a dummy array if no true_rho is available. It won't be used.
        true_rho_data_or_dummy = (
            self.true_rho.data
            if compute_infidelity_flag
            else jnp.empty((dim, dim), dtype=jnp.complex64)
        )

        final_carry, history = _run_tomography_scan(
            initial_params=params,
            initial_opt_state=opt_state,
            true_rho_data=true_rho_data_or_dummy,
            measurement_basis=self.measurement_basis,
            measurement_results=self.measurement_results,
            dim=dim,
            epochs=epochs,
            optimizer=optimizer,
            compute_infidelity=compute_infidelity_flag,
            L1_reg_strength=L1_reg_strength,
        )

        final_params, _ = final_carry

        rho = Qarray.create(_reconstruct_density_matrix(final_params, dim))

        self._result = MLETomographyResult(
            rho=rho,
            params_history=history["params"],
            loss_history=history["loss"],
            grads_history=history["grads"],
            infidelity_history=history["infidelity"]
            if compute_infidelity_flag
            else None,
        )
        return self._result

    def quantum_state_tomography_direct(
        self,
    ) -> Qarray:

        """Perform quantum state tomography using direct inversion.

        This method reconstructs the quantum state from measurement results by 
        directly solving a system of linear equations. The method assumes that
        the measurement basis is complete and the measurement results are 
        noise-free.

        Returns:
            Qarray: Reconstructed quantum state.
        """

    # Compute overlaps of measurement and complete operator bases
        A = jnp.einsum("ijk,ljk->il", self.complete_basis, self.measurement_basis)
        # Solve the linear system to find the coefficients
        coefficients = jnp.linalg.solve(A, self.measurement_results)
        # Reconstruct the density matrix
        rho = jnp.einsum("i, ijk->jk", coefficients, self.complete_basis)

        return Qarray.create(rho)

    def plot_results(self):
        if self._result is None:
            raise ValueError(
                "No results to plot. Run quantum_state_tomography_mle first."
            )

        _, ax = plt.subplots(1, figsize=(5, 4))
        if self._result.infidelity_history is not None:
            ax2 = ax.twinx()

        ax.plot(self._result.loss_history, color="C0")
        ax.set_xlabel("Epoch")
        ax.set_ylabel("$\\mathcal{L}$", color="C0")
        ax.set_yscale("log")

        if self._result.infidelity_history is not None:
            ax2.plot(self._result.infidelity_history, color="C1")
            ax2.set_yscale("log")
            ax2.set_ylabel("$1-\\mathcal{F}$", color="C1")
            plt.grid(False)

        plt.show()

__init__(rho_guess, measurement_basis, measurement_results, complete_basis=None, true_rho=None)

Reconstruct a quantum state from measurement results using quantum state tomography. The tomography can be performed either by direct inversion or by maximum likelihood estimation.

Parameters:

Name	Type	Description	Default
rho_guess	Qarray	The initial guess for the quantum state.	required
measurement_basis	Qarray	The basis in which measurements are performed.	required
measurement_results	ndarray	The results of the measurements.	required
complete_basis	Optional[Qarray]	The complete basis for state	None
true_rho	Optional[Qarray]	The true quantum state, if known.	None

Source code in jaxquantum/core/measurements.py

def __init__(
    self,
    rho_guess: Qarray,
    measurement_basis: Qarray,
    measurement_results: jnp.ndarray,
    complete_basis: Optional[Qarray] = None,
    true_rho: Optional[Qarray] = None,
):
    """
    Reconstruct a quantum state from measurement results using quantum state tomography.
    The tomography can be performed either by direct inversion or by maximum likelihood estimation.

    Args:
        rho_guess (Qarray): The initial guess for the quantum state.
        measurement_basis (Qarray): The basis in which measurements are performed.
        measurement_results (jnp.ndarray): The results of the measurements.
        complete_basis (Optional[Qarray]): The complete basis for state 
        reconstruction used when using direct inversion. 
        Defaults to the measurement basis if not provided.
        true_rho (Optional[Qarray]): The true quantum state, if known.

    """
    self.rho_guess = rho_guess.data
    self.measurement_basis = measurement_basis.data
    self.measurement_results = measurement_results
    self.complete_basis = (
        complete_basis.data
        if (complete_basis is not None)
        else measurement_basis.data
    )
    self.true_rho = true_rho
    self._result = None

quantum_state_tomography_direct()

Perform quantum state tomography using direct inversion.

This method reconstructs the quantum state from measurement results by directly solving a system of linear equations. The method assumes that the measurement basis is complete and the measurement results are noise-free.

Returns:

Name	Type	Description
Qarray	Qarray	Reconstructed quantum state.

Source code in jaxquantum/core/measurements.py

def quantum_state_tomography_direct(
    self,
) -> Qarray:

    """Perform quantum state tomography using direct inversion.

    This method reconstructs the quantum state from measurement results by 
    directly solving a system of linear equations. The method assumes that
    the measurement basis is complete and the measurement results are 
    noise-free.

    Returns:
        Qarray: Reconstructed quantum state.
    """

# Compute overlaps of measurement and complete operator bases
    A = jnp.einsum("ijk,ljk->il", self.complete_basis, self.measurement_basis)
    # Solve the linear system to find the coefficients
    coefficients = jnp.linalg.solve(A, self.measurement_results)
    # Reconstruct the density matrix
    rho = jnp.einsum("i, ijk->jk", coefficients, self.complete_basis)

    return Qarray.create(rho)

quantum_state_tomography_mle(L1_reg_strength=0.0, epochs=10000, lr=0.005)

Perform quantum state tomography using maximum likelihood estimation (MLE).

This method reconstructs the quantum state from measurement results by optimizing a likelihood function using gradient descent. The optimization ensures the resulting density matrix is positive semi-definite with trace 1.

Parameters:

Name	Type	Description	Default
L1_reg_strength	float	Strength of L1	0.0
epochs	int	Number of optimization iterations.	10000
lr	float	Learning rate for the Adam optimizer.	0.005

Returns:

Name	Type	Description
MLETomographyResult	MLETomographyResult	Named tuple containing: - rho: Reconstructed quantum state as Qarray - params_history: List of parameter values during optimization - loss_history: List of loss values during optimization - grads_history: List of gradient values during optimization - infidelity_history: List of infidelities if true_rho was provided, None otherwise

Source code in jaxquantum/core/measurements.py

def quantum_state_tomography_mle(
    self, L1_reg_strength: float = 0.0, epochs: int = 10000, lr: float = 5e-3
) -> MLETomographyResult:
    """Perform quantum state tomography using maximum likelihood 
    estimation (MLE).

    This method reconstructs the quantum state from measurement results 
    by optimizing
    a likelihood function using gradient descent. The optimization 
    ensures the 
    resulting density matrix is positive semi-definite with trace 1.

    Args:
        L1_reg_strength (float, optional): Strength of L1 
        regularization. Defaults to 0.0.
        epochs (int, optional): Number of optimization iterations. 
        Defaults to 10000.
        lr (float, optional): Learning rate for the Adam optimizer. 
        Defaults to 5e-3.

    Returns:
        MLETomographyResult: Named tuple containing:
            - rho: Reconstructed quantum state as Qarray
            - params_history: List of parameter values during optimization
            - loss_history: List of loss values during optimization
            - grads_history: List of gradient values during optimization
            - infidelity_history: List of infidelities if true_rho was 
            provided, None otherwise
    """

    dim = self.rho_guess.shape[0]
    optimizer = optax.adam(lr)

    # Initialize parameters from the initial guess for the density matrix
    params = _parametrize_density_matrix(self.rho_guess, dim)
    opt_state = optimizer.init(params)

    compute_infidelity_flag = self.true_rho is not None

    # Provide a dummy array if no true_rho is available. It won't be used.
    true_rho_data_or_dummy = (
        self.true_rho.data
        if compute_infidelity_flag
        else jnp.empty((dim, dim), dtype=jnp.complex64)
    )

    final_carry, history = _run_tomography_scan(
        initial_params=params,
        initial_opt_state=opt_state,
        true_rho_data=true_rho_data_or_dummy,
        measurement_basis=self.measurement_basis,
        measurement_results=self.measurement_results,
        dim=dim,
        epochs=epochs,
        optimizer=optimizer,
        compute_infidelity=compute_infidelity_flag,
        L1_reg_strength=L1_reg_strength,
    )

    final_params, _ = final_carry

    rho = Qarray.create(_reconstruct_density_matrix(final_params, dim))

    self._result = MLETomographyResult(
        rho=rho,
        params_history=history["params"],
        loss_history=history["loss"],
        grads_history=history["grads"],
        infidelity_history=history["infidelity"]
        if compute_infidelity_flag
        else None,
    )
    return self._result

SolverOptions

Options forwarded to :func:diffrax.diffeqsolve.

Attributes:

Name	Type	Description
solver	AbstractSolver | str	Native Diffrax solver; strings are deprecated.
stepsize_controller	AbstractStepSizeController | str	Native Diffrax controller; strings are deprecated.
stepsize_controller_kwargs	dict[str, Any] | None	Deprecated controller constructor arguments.
saveat	SaveAt | None	Custom save policy; overrides JAXQuantum's save-time handling.
dt0	float | Array | None | Literal['tlist']	Initial step, "tlist" for the first interval, or None for Diffrax's automatic choice.
adjoint	AbstractAdjoint | None	Differentiation strategy. None uses Diffrax's default.
event	Event | None	Native Diffrax termination event.
max_steps	int | None	Maximum solver steps.
throw	bool	Whether unsuccessful solves raise an exception.
progress_meter	bool | Literal['default'] | AbstractProgressMeter | None	None, "default", or a native progress meter. Booleans are deprecated.
solver_state	Any	Solver state used to continue a previous solve.
controller_state	Any	Controller state used to continue a previous solve.
made_jump	bool | Array | None	Previous jump state used when continuing a solve.

saveat=None saves at saveat_tlist (or tlist when omitted). adjoint=None and progress_meter=None preserve Diffrax's defaults. Native Diffrax objects pass through unchanged. Legacy values still work and issue a FutureWarning.

Source code in jaxquantum/core/solvers.py

@struct.dataclass
class SolverOptions:
    """Options forwarded to :func:`diffrax.diffeqsolve`.

    Attributes:
        solver: Native Diffrax solver; strings are deprecated.
        stepsize_controller: Native Diffrax controller; strings are deprecated.
        stepsize_controller_kwargs: Deprecated controller constructor arguments.
        saveat: Custom save policy; overrides JAXQuantum's save-time handling.
        dt0: Initial step, ``"tlist"`` for the first interval, or ``None`` for
            Diffrax's automatic choice.
        adjoint: Differentiation strategy. ``None`` uses Diffrax's default.
        event: Native Diffrax termination event.
        max_steps: Maximum solver steps.
        throw: Whether unsuccessful solves raise an exception.
        progress_meter: ``None``, ``"default"``, or a native progress meter.
            Booleans are deprecated.
        solver_state: Solver state used to continue a previous solve.
        controller_state: Controller state used to continue a previous solve.
        made_jump: Previous jump state used when continuing a solve.

    ``saveat=None`` saves at ``saveat_tlist`` (or ``tlist`` when omitted).
    ``adjoint=None`` and ``progress_meter=None`` preserve Diffrax's defaults.
    Native Diffrax objects pass through unchanged. Legacy values still work and
    issue a ``FutureWarning``.
    """

    progress_meter: bool | Literal["default"] | diffrax.AbstractProgressMeter | None = (
        struct.field(pytree_node=False, default="default")
    )
    solver: diffrax.AbstractSolver | str = struct.field(
        pytree_node=False, default_factory=diffrax.Tsit5
    )
    max_steps: int | None = struct.field(pytree_node=False, default=100_000)
    stepsize_controller: diffrax.AbstractStepSizeController | str = struct.field(
        pytree_node=False, default_factory=_default_stepsize_controller
    )
    stepsize_controller_kwargs: dict[str, Any] | None = struct.field(
        pytree_node=False, default=None
    )
    saveat: diffrax.SaveAt | None = struct.field(pytree_node=False, default=None)
    dt0: float | Array | None | Literal["tlist"] = struct.field(
        pytree_node=False, default="tlist"
    )
    adjoint: diffrax.AbstractAdjoint | None = struct.field(
        pytree_node=False, default=None
    )
    event: diffrax.Event | None = struct.field(pytree_node=False, default=None)
    throw: bool = struct.field(pytree_node=False, default=True)
    solver_state: Any = None
    controller_state: Any = None
    made_jump: bool | Array | None = None

    @classmethod
    def create(
        cls,
        progress_meter: bool = True,
        solver: str = "Tsit5",
        max_steps: int = 100_000,
        stepsize_controller: str = "PIDController",
        stepsize_controller_kwargs: dict[str, Any] | None = None,
    ) -> "SolverOptions":
        """Create options with the deprecated string-based interface."""
        warnings.warn(
            "SolverOptions.create() is deprecated; use SolverOptions with native "
            "objects, such as solver=diffrax.Tsit5() and "
            "progress_meter='default'.",
            FutureWarning,
            stacklevel=2,
        )
        return cls(
            solver=_diffrax_object(solver, diffrax.AbstractSolver),
            stepsize_controller=_diffrax_object(
                stepsize_controller,
                diffrax.AbstractStepSizeController,
                _legacy_controller_kwargs(
                    stepsize_controller, stepsize_controller_kwargs
                ),
            ),
            max_steps=max_steps,
            progress_meter="default" if progress_meter else None,
        )

create(progress_meter=True, solver='Tsit5', max_steps=100000, stepsize_controller='PIDController', stepsize_controller_kwargs=None) classmethod

Create options with the deprecated string-based interface.

Source code in jaxquantum/core/solvers.py

@classmethod
def create(
    cls,
    progress_meter: bool = True,
    solver: str = "Tsit5",
    max_steps: int = 100_000,
    stepsize_controller: str = "PIDController",
    stepsize_controller_kwargs: dict[str, Any] | None = None,
) -> "SolverOptions":
    """Create options with the deprecated string-based interface."""
    warnings.warn(
        "SolverOptions.create() is deprecated; use SolverOptions with native "
        "objects, such as solver=diffrax.Tsit5() and "
        "progress_meter='default'.",
        FutureWarning,
        stacklevel=2,
    )
    return cls(
        solver=_diffrax_object(solver, diffrax.AbstractSolver),
        stepsize_controller=_diffrax_object(
            stepsize_controller,
            diffrax.AbstractStepSizeController,
            _legacy_controller_kwargs(
                stepsize_controller, stepsize_controller_kwargs
            ),
        ),
        max_steps=max_steps,
        progress_meter="default" if progress_meter else None,
    )

SparseBCOOImpl

Bases: QarrayImpl

Sparse implementation using JAX experimental BCOO sparse arrays.

Attributes:

Name	Type	Description
_data	BCOO	The underlying sparse.BCOO array.

Source code in jaxquantum/core/sparse_bcoo.py

@struct.dataclass
class SparseBCOOImpl(QarrayImpl):
    """Sparse implementation using JAX experimental BCOO sparse arrays.

    Attributes:
        _data: The underlying ``sparse.BCOO`` array.
    """

    _data: sparse.BCOO

    PROMOTION_ORDER = 1  # noqa: RUF012 — not a struct field; no annotation intentional

    @classmethod
    def _make(cls, data) -> "SparseBCOOImpl":
        """Construct a ``SparseBCOOImpl`` (no sharding applied).

        SparseBCOO has variable nnz per row and does not shard cleanly, so
        ``_make`` is a pure pass-through. The default-sharding gate lives in
        :meth:`from_data` instead, where users construct sharded Qarrays.
        """
        return cls(_data=data)

    @classmethod
    def from_data(cls, data) -> "SparseBCOOImpl":
        """Wrap *data* in a new ``SparseBCOOImpl``, converting to BCOO if needed.

        Args:
            data: A ``sparse.BCOO`` or array-like input.

        Returns:
            A ``SparseBCOOImpl`` wrapping a BCOO representation of *data*.

        Raises:
            NotImplementedError: If a default sharding is configured. BCOO
                does not partition cleanly across devices; convert to Dense
                or SparseDIA, or call ``jqt.clear_default_sharding()``.
        """
        if SETTINGS["default_sharding"] is not None:
            raise NotImplementedError(
                "Sharding is not supported for SparseBCOO (variable nnz "
                "per shard). Convert to Dense or SparseDIA, or call "
                "`jqt.clear_default_sharding()` before constructing a "
                "SparseBCOO Qarray."
            )
        return cls._make(cls._to_sparse(data))

    def get_data(self) -> Array:
        """Return the underlying BCOO sparse array."""
        return self._data

    def matmul(self, other: QarrayImpl) -> QarrayImpl:
        """Matrix multiply ``self @ other``.

        When *other* is a ``DenseImpl``, JAX's native BCOO @ dense path is
        used (no self-densification).  When *other* is also a
        ``SparseBCOOImpl``, a sparse @ sparse product is performed.

        Args:
            other: Right-hand operand.

        Returns:
            A ``DenseImpl`` (sparse @ dense) or ``SparseBCOOImpl`` (sparse @
            sparse) containing the matrix product.
        """
        if isinstance(other, DenseImpl):
            return DenseImpl._make(self._data @ other._data)
        a, b = self._coerce(other)
        if a is not self:
            return a.matmul(b)
        return SparseBCOOImpl._make(self._data @ b._data)

    def add(self, other: QarrayImpl) -> QarrayImpl:
        """Element-wise addition ``self + other``, coercing types as needed.

        Args:
            other: Right-hand operand.

        Returns:
            A ``SparseBCOOImpl`` (both sparse) or ``DenseImpl`` (mixed) sum.
        """
        a, b = self._coerce(other)
        if a is not self:
            return a.add(b)
        x, y = self._data, b._data
        x, y = self._broadcast_pair(x, y)
        if x.indices.dtype != y.indices.dtype:
            y = sparse.BCOO((y.data, y.indices.astype(x.indices.dtype)), shape=y.shape)
        return SparseBCOOImpl._make(x + y)

    def sub(self, other: QarrayImpl) -> QarrayImpl:
        """Element-wise subtraction ``self - other``, coercing types as needed.

        Args:
            other: Right-hand operand.

        Returns:
            A ``SparseBCOOImpl`` (both sparse) or ``DenseImpl`` (mixed) difference.
        """
        a, b = self._coerce(other)
        if a is not self:
            return a.sub(b)
        x, y = self._data, b._data
        x, y = self._broadcast_pair(x, y)
        if x.indices.dtype != y.indices.dtype:
            y = sparse.BCOO((y.data, y.indices.astype(x.indices.dtype)), shape=y.shape)
        return SparseBCOOImpl._make(x - y)

    @staticmethod
    def _broadcast_pair(x: sparse.BCOO, y: sparse.BCOO):
        shape = jnp.broadcast_shapes(x.shape, y.shape)
        if x.shape != shape:
            dims = range(len(shape) - x.ndim, len(shape))
            x = sparse.bcoo_broadcast_in_dim(
                x, shape=shape, broadcast_dimensions=dims
            )
        if y.shape != shape:
            dims = range(len(shape) - y.ndim, len(shape))
            y = sparse.bcoo_broadcast_in_dim(
                y, shape=shape, broadcast_dimensions=dims
            )
        return x, y

    def mul(self, scalar) -> QarrayImpl:
        """Scalar multiplication.

        Args:
            scalar: Scalar value.

        Returns:
            A ``SparseBCOOImpl`` with each stored value multiplied by *scalar*.
        """
        return SparseBCOOImpl._make(scalar * self._data)

    def dag(self) -> QarrayImpl:
        """Conjugate transpose without densifying.

        Transposes the last two dimensions of the BCOO array and conjugates
        the stored values.

        Returns:
            A ``SparseBCOOImpl`` containing the conjugate transpose.
        """
        ndim = self._data.ndim
        if ndim >= 2:
            permutation = tuple(range(ndim - 2)) + (ndim - 1, ndim - 2)
            transposed_data = sparse.bcoo_transpose(self._data, permutation=permutation)
        else:
            transposed_data = self._data

        conjugated_data = sparse.BCOO(
            (jnp.conj(transposed_data.data), transposed_data.indices),
            shape=transposed_data.shape,
        )
        return SparseBCOOImpl._make(conjugated_data)

    def to_dense(self) -> "DenseImpl":
        """Convert to a ``DenseImpl`` via ``todense()``.

        Returns:
            A ``DenseImpl`` with the same values as this sparse array.
        """
        return DenseImpl._make(self._data.todense())

    @classmethod
    def _to_sparse(cls, data) -> sparse.BCOO:
        """Convert *data* to a ``sparse.BCOO``, returning it unchanged if already sparse.

        Args:
            data: A ``sparse.BCOO`` or array-like.

        Returns:
            A ``sparse.BCOO`` representation of *data*.
        """
        if isinstance(data, sparse.BCOO):
            return data
        return sparse.BCOO.fromdense(data)

    def to_sparse_bcoo(self) -> "SparseBCOOImpl":
        """Return self (already sparse BCOO).

        Returns:
            This ``SparseBCOOImpl`` instance unchanged.
        """
        return self

    def shape(self) -> tuple:
        """Shape of the underlying BCOO array.

        Returns:
            Tuple of dimension sizes.
        """
        return self._data.shape

    def dtype(self):
        """Data type of the underlying BCOO array.

        Returns:
            The dtype of ``_data``.
        """
        return self._data.dtype

    def frobenius_norm(self) -> float:
        """Compute the Frobenius norm directly from stored values.

        Returns:
            The Frobenius norm as a scalar.
        """
        return jnp.sqrt(jnp.sum(jnp.abs(self._data.data) ** 2))

    @classmethod
    def _real(cls, data):
        """Return a BCOO array with only the real parts of the stored values."""
        return sparse.BCOO((jnp.real(data.data), data.indices), shape=data.shape)

    def real(self) -> QarrayImpl:
        """Element-wise real part.

        Returns:
            A ``SparseBCOOImpl`` containing the real parts of stored values.
        """
        return SparseBCOOImpl._make(SparseBCOOImpl._real(self._data))

    @classmethod
    def _imag(cls, data):
        """Return a BCOO array with only the imaginary parts of the stored values."""
        return sparse.BCOO((jnp.imag(data.data), data.indices), shape=data.shape)

    def imag(self) -> QarrayImpl:
        """Element-wise imaginary part.

        Returns:
            A ``SparseBCOOImpl`` containing the imaginary parts of stored values.
        """
        return SparseBCOOImpl._make(SparseBCOOImpl._imag(self._data))

    @classmethod
    def _conj(cls, data):
        """Return a BCOO array with complex-conjugated stored values."""
        return sparse.BCOO((jnp.conj(data.data), data.indices), shape=data.shape)

    def conj(self) -> QarrayImpl:
        """Element-wise complex conjugate.

        Returns:
            A ``SparseBCOOImpl`` containing the complex-conjugated stored values.
        """
        return SparseBCOOImpl._make(SparseBCOOImpl._conj(self._data))

    @classmethod
    def _abs(cls, data):
        """Return a BCOO array with absolute values of stored entries."""
        return sparse.sparsify(jnp.abs)(data)

    def abs(self) -> QarrayImpl:
        """Element-wise absolute value.

        Returns:
            A ``SparseBCOOImpl`` containing the absolute values of stored entries.
        """
        return SparseBCOOImpl._make(SparseBCOOImpl._abs(self._data))

    @classmethod
    def _eye_data(cls, n: int, dtype=None):
        """Create an ``n x n`` identity matrix as a sparse BCOO with O(n) memory.

        Args:
            n: Matrix size.
            dtype: Optional data type.

        Returns:
            A ``sparse.BCOO`` identity matrix of shape ``(n, n)``.
        """
        return sparse.eye(n, dtype=dtype)

    @classmethod
    def _scaled_identity(cls, n: int, scalar, dtype=None) -> SparseBCOOImpl:
        """Create a batched scaled identity without dense broadcasting."""
        scalar = jnp.ones((), dtype=dtype) * (jnp.asarray(scalar) + 0.0j)
        batch_shape = scalar.shape
        diagonal = jnp.arange(n)
        indices = jnp.stack((diagonal, diagonal), axis=-1)
        indices = jnp.broadcast_to(indices, (*batch_shape, n, 2))
        values = jnp.broadcast_to(scalar[..., None], (*batch_shape, n))
        return cls._make(
            sparse.BCOO((values, indices), shape=(*batch_shape, n, n))
        )

    @classmethod
    def can_handle_data(cls, arr) -> bool:
        """Return True when *arr* is a ``sparse.BCOO`` array.

        Args:
            arr: Raw array.

        Returns:
            True if *arr* is a ``sparse.BCOO`` instance.
        """
        return isinstance(arr, sparse.BCOO)

    @classmethod
    def dag_data(cls, arr: sparse.BCOO) -> sparse.BCOO:
        """Conjugate transpose for BCOO sparse arrays without densifying.

        Args:
            arr: A ``sparse.BCOO`` array with ``ndim >= 2``.

        Returns:
            A ``sparse.BCOO`` containing the conjugate transpose.
        """
        ndim = arr.ndim
        permutation = tuple(range(ndim - 2)) + (ndim - 1, ndim - 2)
        transposed = sparse.bcoo_transpose(arr, permutation=permutation)
        return sparse.BCOO(
            (jnp.conj(transposed.data), transposed.indices),
            shape=transposed.shape,
        )

    def trace(self) -> Array:
        """Compute the trace of the last two matrix dimensions without densifying.

        Returns:
            Trace value(s).
        """
        indices = self._data.indices
        values = self._data.data
        ndim = indices.shape[-1]

        is_diag = indices[:, -2] == indices[:, -1]

        if ndim == 2:
            return jnp.sum(values * is_diag)
        else:
            batch_shape = self._data.shape[:-2]
            B = int(jnp.prod(jnp.array(batch_shape)))
            strides = [1]
            for s in reversed(batch_shape[1:]):
                strides.insert(0, strides[0] * s)
            strides = jnp.array(strides, dtype=jnp.int32)
            flat_batch_idx = jnp.sum(indices[:, :-2] * strides, axis=-1)
            result = jnp.zeros(B, dtype=values.dtype).at[flat_batch_idx].add(
                values * is_diag
            )
            return result.reshape(batch_shape)

    def keep_only_diag(self) -> "SparseBCOOImpl":
        """Zero out off-diagonal stored entries without densifying.

        Returns:
            A ``SparseBCOOImpl`` with only diagonal entries non-zero.
        """
        indices = self._data.indices
        values = self._data.data
        is_diag = indices[:, -2] == indices[:, -1]
        new_values = values * is_diag
        return SparseBCOOImpl._make(sparse.BCOO((new_values, indices), shape=self._data.shape))

    def l2_norm_batched(self, bdims: tuple) -> Array:
        """Compute the L2 norm per batch element without densifying.

        Args:
            bdims: Tuple of batch dimension sizes.

        Returns:
            Scalar or array of L2 norms.
        """
        values = self._data.data
        indices = self._data.indices
        n_batch_dims = len(bdims)
        sq = jnp.abs(values) ** 2

        if n_batch_dims == 0:
            return jnp.sqrt(jnp.sum(sq))
        else:
            B = int(jnp.prod(jnp.array(bdims)))
            strides = [1]
            for s in reversed(bdims[1:]):
                strides.insert(0, strides[0] * s)
            strides = jnp.array(strides, dtype=jnp.int32)
            flat_batch_idx = jnp.sum(indices[:, :n_batch_dims] * strides, axis=-1)
            sum_sq = (
                jnp.zeros(B, dtype=jnp.float64)
                .at[flat_batch_idx]
                .add(sq)
            )
            return jnp.sqrt(sum_sq).reshape(bdims)

    def __deepcopy__(self, memo=None):
        return SparseBCOOImpl._make(deepcopy(self._data, memo))

    def tidy_up(self, atol) -> "SparseBCOOImpl":
        """Zero out stored values whose real or imaginary magnitude is below *atol*.

        Args:
            atol: Absolute tolerance threshold.

        Returns:
            A new ``SparseBCOOImpl`` with small values zeroed.
        """
        values = self._data.data
        re = jnp.real(values)
        im = jnp.imag(values)
        new_values = re * (jnp.abs(re) > atol) + 1j * im * (jnp.abs(im) > atol)
        return SparseBCOOImpl._make(
            sparse.BCOO((new_values, self._data.indices), shape=self._data.shape)
        )

    def kron(self, other: "QarrayImpl") -> "QarrayImpl":
        """Kronecker product using ``sparsify(jnp.kron)`` — stays sparse.

        Args:
            other: Right-hand operand.

        Returns:
            A ``SparseBCOOImpl`` containing the Kronecker product when both
            operands are sparse; a ``DenseImpl`` when types differ.
        """
        a, b = self._coerce(other)
        if a is not self:
            return a.kron(b)
        sparse_kron = sparse.sparsify(jnp.kron)
        return SparseBCOOImpl._make(sparse_kron(self._data, b._data))

abs()

Element-wise absolute value.

Returns:

Type	Description
QarrayImpl	A SparseBCOOImpl containing the absolute values of stored entries.

Source code in jaxquantum/core/sparse_bcoo.py

def abs(self) -> QarrayImpl:
    """Element-wise absolute value.

    Returns:
        A ``SparseBCOOImpl`` containing the absolute values of stored entries.
    """
    return SparseBCOOImpl._make(SparseBCOOImpl._abs(self._data))

add(other)

Element-wise addition self + other, coercing types as needed.

Parameters:

Name	Type	Description	Default
other	QarrayImpl	Right-hand operand.	required

Returns:

Type	Description
QarrayImpl	A SparseBCOOImpl (both sparse) or DenseImpl (mixed) sum.

Source code in jaxquantum/core/sparse_bcoo.py

def add(self, other: QarrayImpl) -> QarrayImpl:
    """Element-wise addition ``self + other``, coercing types as needed.

    Args:
        other: Right-hand operand.

    Returns:
        A ``SparseBCOOImpl`` (both sparse) or ``DenseImpl`` (mixed) sum.
    """
    a, b = self._coerce(other)
    if a is not self:
        return a.add(b)
    x, y = self._data, b._data
    x, y = self._broadcast_pair(x, y)
    if x.indices.dtype != y.indices.dtype:
        y = sparse.BCOO((y.data, y.indices.astype(x.indices.dtype)), shape=y.shape)
    return SparseBCOOImpl._make(x + y)

can_handle_data(arr) classmethod

Return True when arr is a sparse.BCOO array.

Parameters:

Name	Type	Description	Default
arr		Raw array.	required

Returns:

Type	Description
bool	True if arr is a sparse.BCOO instance.

Source code in jaxquantum/core/sparse_bcoo.py

@classmethod
def can_handle_data(cls, arr) -> bool:
    """Return True when *arr* is a ``sparse.BCOO`` array.

    Args:
        arr: Raw array.

    Returns:
        True if *arr* is a ``sparse.BCOO`` instance.
    """
    return isinstance(arr, sparse.BCOO)

conj()

Element-wise complex conjugate.

Returns:

Type	Description
QarrayImpl	A SparseBCOOImpl containing the complex-conjugated stored values.

Source code in jaxquantum/core/sparse_bcoo.py

def conj(self) -> QarrayImpl:
    """Element-wise complex conjugate.

    Returns:
        A ``SparseBCOOImpl`` containing the complex-conjugated stored values.
    """
    return SparseBCOOImpl._make(SparseBCOOImpl._conj(self._data))

dag()

Conjugate transpose without densifying.

Transposes the last two dimensions of the BCOO array and conjugates the stored values.

Returns:

Type	Description
QarrayImpl	A SparseBCOOImpl containing the conjugate transpose.

Source code in jaxquantum/core/sparse_bcoo.py

def dag(self) -> QarrayImpl:
    """Conjugate transpose without densifying.

    Transposes the last two dimensions of the BCOO array and conjugates
    the stored values.

    Returns:
        A ``SparseBCOOImpl`` containing the conjugate transpose.
    """
    ndim = self._data.ndim
    if ndim >= 2:
        permutation = tuple(range(ndim - 2)) + (ndim - 1, ndim - 2)
        transposed_data = sparse.bcoo_transpose(self._data, permutation=permutation)
    else:
        transposed_data = self._data

    conjugated_data = sparse.BCOO(
        (jnp.conj(transposed_data.data), transposed_data.indices),
        shape=transposed_data.shape,
    )
    return SparseBCOOImpl._make(conjugated_data)

dag_data(arr) classmethod

Conjugate transpose for BCOO sparse arrays without densifying.

Parameters:

Name	Type	Description	Default
arr	BCOO	A sparse.BCOO array with ndim >= 2.	required

Returns:

Type	Description
BCOO	A sparse.BCOO containing the conjugate transpose.

Source code in jaxquantum/core/sparse_bcoo.py

@classmethod
def dag_data(cls, arr: sparse.BCOO) -> sparse.BCOO:
    """Conjugate transpose for BCOO sparse arrays without densifying.

    Args:
        arr: A ``sparse.BCOO`` array with ``ndim >= 2``.

    Returns:
        A ``sparse.BCOO`` containing the conjugate transpose.
    """
    ndim = arr.ndim
    permutation = tuple(range(ndim - 2)) + (ndim - 1, ndim - 2)
    transposed = sparse.bcoo_transpose(arr, permutation=permutation)
    return sparse.BCOO(
        (jnp.conj(transposed.data), transposed.indices),
        shape=transposed.shape,
    )

dtype()

Data type of the underlying BCOO array.

Returns:

Type	Description
	The dtype of _data.

Source code in jaxquantum/core/sparse_bcoo.py

def dtype(self):
    """Data type of the underlying BCOO array.

    Returns:
        The dtype of ``_data``.
    """
    return self._data.dtype

frobenius_norm()

Compute the Frobenius norm directly from stored values.

Returns:

Type	Description
float	The Frobenius norm as a scalar.

Source code in jaxquantum/core/sparse_bcoo.py

def frobenius_norm(self) -> float:
    """Compute the Frobenius norm directly from stored values.

    Returns:
        The Frobenius norm as a scalar.
    """
    return jnp.sqrt(jnp.sum(jnp.abs(self._data.data) ** 2))

from_data(data) classmethod

Wrap data in a new SparseBCOOImpl, converting to BCOO if needed.

Parameters:

Name	Type	Description	Default
data		A sparse.BCOO or array-like input.	required

Returns:

Type	Description
'SparseBCOOImpl'	A SparseBCOOImpl wrapping a BCOO representation of data.

Raises:

Type	Description
NotImplementedError	If a default sharding is configured. BCOO does not partition cleanly across devices; convert to Dense or SparseDIA, or call jqt.clear_default_sharding().

Source code in jaxquantum/core/sparse_bcoo.py

@classmethod
def from_data(cls, data) -> "SparseBCOOImpl":
    """Wrap *data* in a new ``SparseBCOOImpl``, converting to BCOO if needed.

    Args:
        data: A ``sparse.BCOO`` or array-like input.

    Returns:
        A ``SparseBCOOImpl`` wrapping a BCOO representation of *data*.

    Raises:
        NotImplementedError: If a default sharding is configured. BCOO
            does not partition cleanly across devices; convert to Dense
            or SparseDIA, or call ``jqt.clear_default_sharding()``.
    """
    if SETTINGS["default_sharding"] is not None:
        raise NotImplementedError(
            "Sharding is not supported for SparseBCOO (variable nnz "
            "per shard). Convert to Dense or SparseDIA, or call "
            "`jqt.clear_default_sharding()` before constructing a "
            "SparseBCOO Qarray."
        )
    return cls._make(cls._to_sparse(data))

get_data()

Return the underlying BCOO sparse array.

Source code in jaxquantum/core/sparse_bcoo.py

def get_data(self) -> Array:
    """Return the underlying BCOO sparse array."""
    return self._data

imag()

Element-wise imaginary part.

Returns:

Type	Description
QarrayImpl	A SparseBCOOImpl containing the imaginary parts of stored values.

Source code in jaxquantum/core/sparse_bcoo.py

def imag(self) -> QarrayImpl:
    """Element-wise imaginary part.

    Returns:
        A ``SparseBCOOImpl`` containing the imaginary parts of stored values.
    """
    return SparseBCOOImpl._make(SparseBCOOImpl._imag(self._data))

keep_only_diag()

Zero out off-diagonal stored entries without densifying.

Returns:

Type	Description
'SparseBCOOImpl'	A SparseBCOOImpl with only diagonal entries non-zero.

Source code in jaxquantum/core/sparse_bcoo.py

def keep_only_diag(self) -> "SparseBCOOImpl":
    """Zero out off-diagonal stored entries without densifying.

    Returns:
        A ``SparseBCOOImpl`` with only diagonal entries non-zero.
    """
    indices = self._data.indices
    values = self._data.data
    is_diag = indices[:, -2] == indices[:, -1]
    new_values = values * is_diag
    return SparseBCOOImpl._make(sparse.BCOO((new_values, indices), shape=self._data.shape))

kron(other)

Kronecker product using sparsify(jnp.kron) — stays sparse.

Parameters:

Name	Type	Description	Default
other	'QarrayImpl'	Right-hand operand.	required

Returns:

Type	Description
'QarrayImpl'	A SparseBCOOImpl containing the Kronecker product when both
'QarrayImpl'	operands are sparse; a DenseImpl when types differ.

Source code in jaxquantum/core/sparse_bcoo.py

def kron(self, other: "QarrayImpl") -> "QarrayImpl":
    """Kronecker product using ``sparsify(jnp.kron)`` — stays sparse.

    Args:
        other: Right-hand operand.

    Returns:
        A ``SparseBCOOImpl`` containing the Kronecker product when both
        operands are sparse; a ``DenseImpl`` when types differ.
    """
    a, b = self._coerce(other)
    if a is not self:
        return a.kron(b)
    sparse_kron = sparse.sparsify(jnp.kron)
    return SparseBCOOImpl._make(sparse_kron(self._data, b._data))

l2_norm_batched(bdims)

Compute the L2 norm per batch element without densifying.

Parameters:

Name	Type	Description	Default
bdims	tuple	Tuple of batch dimension sizes.	required

Returns:

Type	Description
Array	Scalar or array of L2 norms.

Source code in jaxquantum/core/sparse_bcoo.py

def l2_norm_batched(self, bdims: tuple) -> Array:
    """Compute the L2 norm per batch element without densifying.

    Args:
        bdims: Tuple of batch dimension sizes.

    Returns:
        Scalar or array of L2 norms.
    """
    values = self._data.data
    indices = self._data.indices
    n_batch_dims = len(bdims)
    sq = jnp.abs(values) ** 2

    if n_batch_dims == 0:
        return jnp.sqrt(jnp.sum(sq))
    else:
        B = int(jnp.prod(jnp.array(bdims)))
        strides = [1]
        for s in reversed(bdims[1:]):
            strides.insert(0, strides[0] * s)
        strides = jnp.array(strides, dtype=jnp.int32)
        flat_batch_idx = jnp.sum(indices[:, :n_batch_dims] * strides, axis=-1)
        sum_sq = (
            jnp.zeros(B, dtype=jnp.float64)
            .at[flat_batch_idx]
            .add(sq)
        )
        return jnp.sqrt(sum_sq).reshape(bdims)

matmul(other)

Matrix multiply self @ other.

When other is a DenseImpl, JAX's native BCOO @ dense path is used (no self-densification). When other is also a SparseBCOOImpl, a sparse @ sparse product is performed.

Parameters:

Name	Type	Description	Default
other	QarrayImpl	Right-hand operand.	required

Returns:

Type	Description
QarrayImpl	A DenseImpl (sparse @ dense) or SparseBCOOImpl (sparse @
QarrayImpl	sparse) containing the matrix product.

Source code in jaxquantum/core/sparse_bcoo.py

def matmul(self, other: QarrayImpl) -> QarrayImpl:
    """Matrix multiply ``self @ other``.

    When *other* is a ``DenseImpl``, JAX's native BCOO @ dense path is
    used (no self-densification).  When *other* is also a
    ``SparseBCOOImpl``, a sparse @ sparse product is performed.

    Args:
        other: Right-hand operand.

    Returns:
        A ``DenseImpl`` (sparse @ dense) or ``SparseBCOOImpl`` (sparse @
        sparse) containing the matrix product.
    """
    if isinstance(other, DenseImpl):
        return DenseImpl._make(self._data @ other._data)
    a, b = self._coerce(other)
    if a is not self:
        return a.matmul(b)
    return SparseBCOOImpl._make(self._data @ b._data)

mul(scalar)

Scalar multiplication.

Parameters:

Name	Type	Description	Default
scalar		Scalar value.	required

Returns:

Type	Description
QarrayImpl	A SparseBCOOImpl with each stored value multiplied by scalar.

Source code in jaxquantum/core/sparse_bcoo.py

def mul(self, scalar) -> QarrayImpl:
    """Scalar multiplication.

    Args:
        scalar: Scalar value.

    Returns:
        A ``SparseBCOOImpl`` with each stored value multiplied by *scalar*.
    """
    return SparseBCOOImpl._make(scalar * self._data)

real()

Element-wise real part.

Returns:

Type	Description
QarrayImpl	A SparseBCOOImpl containing the real parts of stored values.

Source code in jaxquantum/core/sparse_bcoo.py

def real(self) -> QarrayImpl:
    """Element-wise real part.

    Returns:
        A ``SparseBCOOImpl`` containing the real parts of stored values.
    """
    return SparseBCOOImpl._make(SparseBCOOImpl._real(self._data))

shape()

Shape of the underlying BCOO array.

Returns:

Type	Description
tuple	Tuple of dimension sizes.

Source code in jaxquantum/core/sparse_bcoo.py

def shape(self) -> tuple:
    """Shape of the underlying BCOO array.

    Returns:
        Tuple of dimension sizes.
    """
    return self._data.shape

sub(other)

Element-wise subtraction self - other, coercing types as needed.

Parameters:

Name	Type	Description	Default
other	QarrayImpl	Right-hand operand.	required

Returns:

Type	Description
QarrayImpl	A SparseBCOOImpl (both sparse) or DenseImpl (mixed) difference.

Source code in jaxquantum/core/sparse_bcoo.py

def sub(self, other: QarrayImpl) -> QarrayImpl:
    """Element-wise subtraction ``self - other``, coercing types as needed.

    Args:
        other: Right-hand operand.

    Returns:
        A ``SparseBCOOImpl`` (both sparse) or ``DenseImpl`` (mixed) difference.
    """
    a, b = self._coerce(other)
    if a is not self:
        return a.sub(b)
    x, y = self._data, b._data
    x, y = self._broadcast_pair(x, y)
    if x.indices.dtype != y.indices.dtype:
        y = sparse.BCOO((y.data, y.indices.astype(x.indices.dtype)), shape=y.shape)
    return SparseBCOOImpl._make(x - y)

tidy_up(atol)

Zero out stored values whose real or imaginary magnitude is below atol.

Parameters:

Name	Type	Description	Default
atol		Absolute tolerance threshold.	required

Returns:

Type	Description
'SparseBCOOImpl'	A new SparseBCOOImpl with small values zeroed.

Source code in jaxquantum/core/sparse_bcoo.py

def tidy_up(self, atol) -> "SparseBCOOImpl":
    """Zero out stored values whose real or imaginary magnitude is below *atol*.

    Args:
        atol: Absolute tolerance threshold.

    Returns:
        A new ``SparseBCOOImpl`` with small values zeroed.
    """
    values = self._data.data
    re = jnp.real(values)
    im = jnp.imag(values)
    new_values = re * (jnp.abs(re) > atol) + 1j * im * (jnp.abs(im) > atol)
    return SparseBCOOImpl._make(
        sparse.BCOO((new_values, self._data.indices), shape=self._data.shape)
    )

to_dense()

Convert to a DenseImpl via todense().

Returns:

Type	Description
'DenseImpl'	A DenseImpl with the same values as this sparse array.

Source code in jaxquantum/core/sparse_bcoo.py

def to_dense(self) -> "DenseImpl":
    """Convert to a ``DenseImpl`` via ``todense()``.

    Returns:
        A ``DenseImpl`` with the same values as this sparse array.
    """
    return DenseImpl._make(self._data.todense())

to_sparse_bcoo()

Return self (already sparse BCOO).

Returns:

Type	Description
'SparseBCOOImpl'	This SparseBCOOImpl instance unchanged.

Source code in jaxquantum/core/sparse_bcoo.py

def to_sparse_bcoo(self) -> "SparseBCOOImpl":
    """Return self (already sparse BCOO).

    Returns:
        This ``SparseBCOOImpl`` instance unchanged.
    """
    return self

trace()

Compute the trace of the last two matrix dimensions without densifying.

Returns:

Type	Description
Array	Trace value(s).

Source code in jaxquantum/core/sparse_bcoo.py

def trace(self) -> Array:
    """Compute the trace of the last two matrix dimensions without densifying.

    Returns:
        Trace value(s).
    """
    indices = self._data.indices
    values = self._data.data
    ndim = indices.shape[-1]

    is_diag = indices[:, -2] == indices[:, -1]

    if ndim == 2:
        return jnp.sum(values * is_diag)
    else:
        batch_shape = self._data.shape[:-2]
        B = int(jnp.prod(jnp.array(batch_shape)))
        strides = [1]
        for s in reversed(batch_shape[1:]):
            strides.insert(0, strides[0] * s)
        strides = jnp.array(strides, dtype=jnp.int32)
        flat_batch_idx = jnp.sum(indices[:, :-2] * strides, axis=-1)
        result = jnp.zeros(B, dtype=values.dtype).at[flat_batch_idx].add(
            values * is_diag
        )
        return result.reshape(batch_shape)

SparseDiaData

Lightweight pytree-compatible container for sparse-diagonal raw data.

Returned by SparseDiaImpl.get_data() and consumed by SparseDiaImpl.from_data(). Registered as a JAX pytree via Flax's @struct.dataclass; offsets is not a pytree leaf (it is static compile-time metadata).

Attributes:

Name	Type	Description
offsets	tuple	Static tuple of diagonal offsets (pytree_node=False).
diags	Array	JAX array of shape (*batch, n_diags, n) containing the padded diagonal values.

Source code in jaxquantum/core/sparse_dia.py

@struct.dataclass
class SparseDiaData:
    """Lightweight pytree-compatible container for sparse-diagonal raw data.

    Returned by ``SparseDiaImpl.get_data()`` and consumed by
    ``SparseDiaImpl.from_data()``.  Registered as a JAX pytree via Flax's
    ``@struct.dataclass``; ``offsets`` is *not* a pytree leaf (it is static
    compile-time metadata).

    Attributes:
        offsets: Static tuple of diagonal offsets (pytree_node=False).
        diags:   JAX array of shape (*batch, n_diags, n) containing the
                 padded diagonal values.
    """

    offsets: tuple = struct.field(pytree_node=False)
    diags: Array

    # Class-level marker (not a dataclass field — no type annotation).
    # Used by DenseImpl.can_handle_data to exclude SparseDiaData without
    # a direct import (which would be circular).
    _is_sparse_dia = True

    @property
    def shape(self) -> tuple:
        """Shape of the represented square matrix (*batch, n, n)."""
        n = self.diags.shape[-1]
        return (*self.diags.shape[:-2], n, n)

    @property
    def dtype(self):
        """Dtype of the stored diagonal values."""
        return self.diags.dtype

    def __mul__(self, scalar):
        return SparseDiaData(offsets=self.offsets, diags=self.diags * scalar)

    def __rmul__(self, scalar):
        return SparseDiaData(offsets=self.offsets, diags=scalar * self.diags)

    def __getitem__(self, index):
        """Index into the batch dimension(s), preserving offsets."""
        return SparseDiaData(offsets=self.offsets, diags=self.diags[index])

    def __len__(self):
        """Number of elements along the leading batch dimension."""
        return self.shape[0]

    def reshape(self, *new_shape):
        """Reshape batch dimensions while preserving diagonal structure.

        ``new_shape`` must end with ``(N, N)`` (the matrix dims are unchanged).
        Only the leading batch dims are reshaped.
        """
        if len(new_shape) == 1 and isinstance(new_shape[0], (tuple, list)):
            new_shape = new_shape[0]
        new_batch = new_shape[:-2]
        n = self.diags.shape[-1]
        new_diags = self.diags.reshape(*new_batch, len(self.offsets), n)
        return SparseDiaData(offsets=self.offsets, diags=new_diags)

    def __matmul__(self, other):
        """SparseDIA @ dense → dense (used by mesolve ODE RHS)."""
        # _sparsedia_matmul_dense is defined later in this module; Python
        # resolves the name at call time so forward reference is fine.
        return _sparsedia_matmul_dense(self.offsets, self.diags, other)

    def __rmatmul__(self, other):
        """dense @ SparseDIA → dense (used by mesolve ODE RHS)."""
        return _sparsedia_rmatmul_dense(other, self.offsets, self.diags)

dtype property

Dtype of the stored diagonal values.

shape property

Shape of the represented square matrix (*batch, n, n).

__getitem__(index)

Index into the batch dimension(s), preserving offsets.

Source code in jaxquantum/core/sparse_dia.py

def __getitem__(self, index):
    """Index into the batch dimension(s), preserving offsets."""
    return SparseDiaData(offsets=self.offsets, diags=self.diags[index])

__len__()

Number of elements along the leading batch dimension.

Source code in jaxquantum/core/sparse_dia.py

def __len__(self):
    """Number of elements along the leading batch dimension."""
    return self.shape[0]

__matmul__(other)

SparseDIA @ dense → dense (used by mesolve ODE RHS).

Source code in jaxquantum/core/sparse_dia.py

def __matmul__(self, other):
    """SparseDIA @ dense → dense (used by mesolve ODE RHS)."""
    # _sparsedia_matmul_dense is defined later in this module; Python
    # resolves the name at call time so forward reference is fine.
    return _sparsedia_matmul_dense(self.offsets, self.diags, other)

__rmatmul__(other)

dense @ SparseDIA → dense (used by mesolve ODE RHS).

Source code in jaxquantum/core/sparse_dia.py

def __rmatmul__(self, other):
    """dense @ SparseDIA → dense (used by mesolve ODE RHS)."""
    return _sparsedia_rmatmul_dense(other, self.offsets, self.diags)

reshape(*new_shape)

Reshape batch dimensions while preserving diagonal structure.

new_shape must end with (N, N) (the matrix dims are unchanged). Only the leading batch dims are reshaped.

Source code in jaxquantum/core/sparse_dia.py

def reshape(self, *new_shape):
    """Reshape batch dimensions while preserving diagonal structure.

    ``new_shape`` must end with ``(N, N)`` (the matrix dims are unchanged).
    Only the leading batch dims are reshaped.
    """
    if len(new_shape) == 1 and isinstance(new_shape[0], (tuple, list)):
        new_shape = new_shape[0]
    new_batch = new_shape[:-2]
    n = self.diags.shape[-1]
    new_diags = self.diags.reshape(*new_batch, len(self.offsets), n)
    return SparseDiaData(offsets=self.offsets, diags=new_diags)

SparseDiaImpl

Bases: QarrayImpl

Sparse-diagonal backend storing only diagonal values.

Data layout::

_offsets  : tuple[int, ...]          — static (pytree_node=False)
_diags    : Array[*batch, n_diags, n] — JAX-traced values

For offset k, the convention is: * k ≥ 0 : valid data at _diags[..., i, k:], zeros at [0:k] * k < 0 : valid data at _diags[..., i, :n+k], zeros at [n+k:]

In both cases: A[row, row+k] = _diags[..., i, row+k]

Source code in jaxquantum/core/sparse_dia.py

@struct.dataclass
class SparseDiaImpl(QarrayImpl):
    """Sparse-diagonal backend storing only diagonal values.

    Data layout::

        _offsets  : tuple[int, ...]          — static (pytree_node=False)
        _diags    : Array[*batch, n_diags, n] — JAX-traced values

    For offset k, the convention is:
        * k ≥ 0 : valid data at ``_diags[..., i, k:]``, zeros at ``[0:k]``
        * k < 0 : valid data at ``_diags[..., i, :n+k]``, zeros at ``[n+k:]``

    In both cases: ``A[row, row+k] = _diags[..., i, row+k]``
    """

    _offsets: tuple = struct.field(pytree_node=False)
    _diags: Array

    PROMOTION_ORDER = 0  # noqa: RUF012 — not a struct field

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    # NOTE (potential future optimization):
    # `_diags` has shape (*batch, n_diags, n) where n_diags is typically 1-5.
    # Sharding the n axis aligns reasonably with dense ops, but for tightly
    # coupled SparseDIA-only chains it may be cheaper to *always replicate*
    # `_diags` (small relative to dense (n, n) storage) and rely on the dense
    # state's sharding to drive parallelism. Benchmark before committing — this
    # is a behavior change for users who want SparseDIA sharding intentionally.
    @classmethod
    def _make(cls, offsets: tuple, diags: Array) -> "SparseDiaImpl":
        """Construct a ``SparseDiaImpl``, applying the configured default sharding.

        All internal construction sites route through this so every Qarray
        (including intermediates produced by ``add``, ``matmul``, ``kron``,
        etc.) satisfies the user's sharding invariant.
        """
        return cls(_offsets=offsets, _diags=_maybe_shard(diags))

    @classmethod
    def from_data(cls, data) -> "SparseDiaImpl":
        """Wrap *data* in a new ``SparseDiaImpl``.

        Accepts either a :class:`SparseDiaData` container (direct wrap) or
        a dense array-like (auto-detect non-zero diagonals via numpy, safe
        to call before JIT).

        Args:
            data: A :class:`SparseDiaData` or dense array of shape
                (*batch, n, n).

        Returns:
            A new ``SparseDiaImpl`` instance.
        """
        if isinstance(data, SparseDiaData):
            return cls._make(data.offsets, data.diags)
        offsets, diags_np = _dense_to_sparsedia(np.asarray(data))
        return cls._make(offsets, jnp.array(diags_np))

    @classmethod
    def from_diags(cls, offsets: tuple, diags: Array) -> "SparseDiaImpl":
        """Directly construct from sorted offsets and padded diagonal array.

        This is the preferred factory when diagonal structure is known in
        advance (e.g., when building ``destroy`` or ``create`` operators).

        Args:
            offsets: Tuple of integer diagonal offsets (need not be sorted;
                will be sorted internally).
            diags:   JAX array of shape (*batch, n_diags, n) with padded
                     diagonal values matching *offsets*.

        Returns:
            A new ``SparseDiaImpl`` instance.
        """
        return cls._make(tuple(sorted(offsets)), diags)

    # ------------------------------------------------------------------
    # QarrayImpl abstract methods
    # ------------------------------------------------------------------

    def get_data(self) -> SparseDiaData:
        """Return a :class:`SparseDiaData` container with the raw diagonal data."""
        return SparseDiaData(offsets=self._offsets, diags=self._diags)

    def shape(self) -> tuple:
        """Shape of the represented square matrix (including batch dims)."""
        n = self._diags.shape[-1]
        return (*self._diags.shape[:-2], n, n)

    def dtype(self):
        """Dtype of the stored diagonal values."""
        return self._diags.dtype

    def __deepcopy__(self, memo=None):
        return SparseDiaImpl._make(deepcopy(self._offsets), self._diags)

    # ------------------------------------------------------------------
    # Arithmetic
    # ------------------------------------------------------------------

    def mul(self, scalar) -> "SparseDiaImpl":
        """Scalar multiplication — scales all diagonal values."""
        return SparseDiaImpl._make(self._offsets, scalar * self._diags)

    def neg(self) -> "SparseDiaImpl":
        """Negation."""
        return SparseDiaImpl._make(self._offsets, -self._diags)

    def add(self, other: QarrayImpl) -> QarrayImpl:
        """Element-wise addition.

        SparseDIA + SparseDIA stays SparseDIA (union of offsets, static).
        Otherwise coerces to the higher-order type.
        """
        if isinstance(other, SparseDiaImpl):
            return _sparsedia_add(self, other)
        a, b = self._coerce(other)
        if a is not self:
            return a.add(b)
        return a.add(b)

    def sub(self, other: QarrayImpl) -> QarrayImpl:
        """Element-wise subtraction."""
        if isinstance(other, SparseDiaImpl):
            return _sparsedia_add(self, other, subtract=True)
        a, b = self._coerce(other)
        if a is not self:
            return a.sub(b)
        return a.sub(b)

    def matmul(self, other: QarrayImpl) -> QarrayImpl:
        """Matrix multiplication.

        * SparseDIA @ SparseDIA → SparseDIA  (O(d₁·d₂·n))
        * SparseDIA @ Dense    → Dense       (O(d·n²), no densification of self)
        * Others               → coerce then delegate
        """
        if isinstance(other, DenseImpl):
            return DenseImpl._make(_sparsedia_matmul_dense(
                self._offsets, self._diags, other._data
            ))
        if isinstance(other, SparseDiaImpl):
            offsets, diags = _sparsedia_matmul_sparsedia(
                self._offsets, self._diags,
                other._offsets, other._diags,
            )
            return SparseDiaImpl._make(offsets, diags)
        a, b = self._coerce(other)
        if a is not self:
            return a.matmul(b)
        return a.matmul(b)

    def dag(self) -> "SparseDiaImpl":
        """Conjugate transpose without densification.

        Negates every offset and rearranges the stored values so that the
        padding convention remains consistent.
        """
        new_offsets = tuple(-k for k in self._offsets)
        new_diags = jnp.zeros_like(self._diags)
        for i, k in enumerate(self._offsets):
            s = _dia_slice(k)    # valid data slice for offset k
            sm = _dia_slice(-k)  # valid data slice for offset -k (the new position)
            new_diags = new_diags.at[..., i, sm].set(jnp.conj(self._diags[..., i, s]))
        return SparseDiaImpl._make(new_offsets, new_diags)

    def kron(self, other: QarrayImpl) -> QarrayImpl:
        """Kronecker product.

        SparseDIA ⊗ SparseDIA stays SparseDIA: output offset for pair
        (kA, kB) is ``kA * m + kB`` where m = dim(B).  Fully vectorised —
        no loops at JAX level.
        """
        if isinstance(other, SparseDiaImpl):
            return _sparsedia_kron(self, other)
        a, b = self._coerce(other)
        if a is not self:
            return a.kron(b)
        return a.kron(b)

    def tidy_up(self, atol) -> "SparseDiaImpl":
        """Zero diagonal values whose magnitude is below *atol*."""
        diags = self._diags
        real_part = jnp.where(jnp.abs(jnp.real(diags)) < atol, 0.0, jnp.real(diags))
        if jnp.issubdtype(diags.dtype, jnp.complexfloating):
            imag_part = jnp.where(jnp.abs(jnp.imag(diags)) < atol, 0.0, jnp.imag(diags))
            new_diags = (real_part + 1j * imag_part).astype(diags.dtype)
        else:
            new_diags = real_part.astype(diags.dtype)
        return SparseDiaImpl._make(self._offsets, new_diags)

    # ------------------------------------------------------------------
    # Conversions
    # ------------------------------------------------------------------

    def to_dense(self) -> "DenseImpl":
        """Convert to a ``DenseImpl`` by summing diagonal contributions."""
        n = self._diags.shape[-1]
        batch_shape = self._diags.shape[:-2]
        result = jnp.zeros((*batch_shape, n, n), dtype=self._diags.dtype)
        for i, k in enumerate(self._offsets):
            s = _dia_slice(k)
            length = n - abs(k)
            if length <= 0:
                continue
            vals = self._diags[..., i, s]
            row_idx = jnp.arange(length) + max(-k, 0)
            col_idx = row_idx + k
            result = result.at[..., row_idx, col_idx].set(vals)
        return DenseImpl._make(result)

    def to_sparse_bcoo(self) -> "SparseBCOOImpl":
        """Convert to a ``SparseBCOOImpl`` (BCOO) via dense."""
        return self.to_dense().to_sparse_bcoo()

    def to_sparse_dia(self) -> "SparseDiaImpl":
        """Return self (already SparseDIA)."""
        return self

    # ------------------------------------------------------------------
    # Class-method interface
    # ------------------------------------------------------------------

    @classmethod
    def _eye_data(cls, n: int, dtype=None):
        """Return an n×n identity as a dense JAX array."""
        return jnp.eye(n, dtype=dtype)

    @classmethod
    def _scaled_identity(cls, n: int, scalar, dtype=None) -> "SparseDiaImpl":
        """Create a batched scaled identity without dense storage."""
        scalar = jnp.ones((), dtype=dtype) * (jnp.asarray(scalar) + 0.0j)
        diags = jnp.broadcast_to(
            scalar[..., None, None],
            (*scalar.shape, 1, n),
        )
        return cls._make((0,), diags)

    @classmethod
    def can_handle_data(cls, arr) -> bool:
        """Return True only for :class:`SparseDiaData` objects."""
        return isinstance(arr, SparseDiaData)

    @classmethod
    def dag_data(cls, arr: SparseDiaData) -> SparseDiaData:
        """Conjugate transpose of raw :class:`SparseDiaData` without densification."""
        impl = SparseDiaImpl._make(arr.offsets, arr.diags)
        result = impl.dag()
        return result.get_data()

    # ------------------------------------------------------------------
    # Extra sparse-native methods (no densification)
    # ------------------------------------------------------------------

    def trace(self):
        """Compute trace directly from the main diagonal (offset 0).

        Returns:
            Scalar trace (sum of main diagonal values).
        """
        if 0 in self._offsets:
            i = self._offsets.index(0)
            return jnp.sum(self._diags[..., i, :], axis=-1)
        return jnp.zeros(self._diags.shape[:-2], dtype=self._diags.dtype)

    def frobenius_norm(self):
        """Frobenius norm computed directly from stored diagonal values."""
        return jnp.sqrt(jnp.sum(jnp.abs(self._diags) ** 2))

    def real(self) -> "SparseDiaImpl":
        """Element-wise real part of stored values."""
        return SparseDiaImpl._make(
            self._offsets,
            jnp.real(self._diags).astype(self._diags.dtype),
        )

    def imag(self) -> "SparseDiaImpl":
        """Element-wise imaginary part of stored values."""
        return SparseDiaImpl._make(
            self._offsets,
            jnp.imag(self._diags).astype(self._diags.dtype),
        )

    def conj(self) -> "SparseDiaImpl":
        """Element-wise complex conjugate of stored values."""
        return SparseDiaImpl._make(self._offsets, jnp.conj(self._diags))

    def powm(self, n: int) -> "SparseDiaImpl":
        """Integer matrix power staying SparseDIA via binary exponentiation.

        Uses O(log n) SparseDIA @ SparseDIA multiplications rather than
        densifying.  A^0 returns the identity operator.

        Args:
            n: Non-negative integer exponent.

        Returns:
            A ``SparseDiaImpl`` equal to this matrix raised to the *n*-th power.

        Raises:
            ValueError: If *n* is negative.
        """
        if n < 0:
            raise ValueError("powm requires n >= 0")
        if n == 0:
            size = self._diags.shape[-1]
            eye_diags = jnp.ones((*self._diags.shape[:-2], 1, size), dtype=self._diags.dtype)
            return SparseDiaImpl._make((0,), eye_diags)
        if n == 1:
            return self
        half = self.powm(n // 2)
        squared = half.matmul(half)  # SparseDIA @ SparseDIA → SparseDIA
        return squared if n % 2 == 0 else self.matmul(squared)

add(other)

Element-wise addition.

SparseDIA + SparseDIA stays SparseDIA (union of offsets, static). Otherwise coerces to the higher-order type.

Source code in jaxquantum/core/sparse_dia.py

def add(self, other: QarrayImpl) -> QarrayImpl:
    """Element-wise addition.

    SparseDIA + SparseDIA stays SparseDIA (union of offsets, static).
    Otherwise coerces to the higher-order type.
    """
    if isinstance(other, SparseDiaImpl):
        return _sparsedia_add(self, other)
    a, b = self._coerce(other)
    if a is not self:
        return a.add(b)
    return a.add(b)

can_handle_data(arr) classmethod

Return True only for :class:SparseDiaData objects.

Source code in jaxquantum/core/sparse_dia.py

@classmethod
def can_handle_data(cls, arr) -> bool:
    """Return True only for :class:`SparseDiaData` objects."""
    return isinstance(arr, SparseDiaData)

conj()

Element-wise complex conjugate of stored values.

Source code in jaxquantum/core/sparse_dia.py

def conj(self) -> "SparseDiaImpl":
    """Element-wise complex conjugate of stored values."""
    return SparseDiaImpl._make(self._offsets, jnp.conj(self._diags))

dag()

Conjugate transpose without densification.

Negates every offset and rearranges the stored values so that the padding convention remains consistent.

Source code in jaxquantum/core/sparse_dia.py

def dag(self) -> "SparseDiaImpl":
    """Conjugate transpose without densification.

    Negates every offset and rearranges the stored values so that the
    padding convention remains consistent.
    """
    new_offsets = tuple(-k for k in self._offsets)
    new_diags = jnp.zeros_like(self._diags)
    for i, k in enumerate(self._offsets):
        s = _dia_slice(k)    # valid data slice for offset k
        sm = _dia_slice(-k)  # valid data slice for offset -k (the new position)
        new_diags = new_diags.at[..., i, sm].set(jnp.conj(self._diags[..., i, s]))
    return SparseDiaImpl._make(new_offsets, new_diags)

dag_data(arr) classmethod

Conjugate transpose of raw :class:SparseDiaData without densification.

Source code in jaxquantum/core/sparse_dia.py

@classmethod
def dag_data(cls, arr: SparseDiaData) -> SparseDiaData:
    """Conjugate transpose of raw :class:`SparseDiaData` without densification."""
    impl = SparseDiaImpl._make(arr.offsets, arr.diags)
    result = impl.dag()
    return result.get_data()

dtype()

Dtype of the stored diagonal values.

Source code in jaxquantum/core/sparse_dia.py

def dtype(self):
    """Dtype of the stored diagonal values."""
    return self._diags.dtype

frobenius_norm()

Frobenius norm computed directly from stored diagonal values.

Source code in jaxquantum/core/sparse_dia.py

def frobenius_norm(self):
    """Frobenius norm computed directly from stored diagonal values."""
    return jnp.sqrt(jnp.sum(jnp.abs(self._diags) ** 2))

from_data(data) classmethod

Wrap data in a new SparseDiaImpl.

Accepts either a :class:SparseDiaData container (direct wrap) or a dense array-like (auto-detect non-zero diagonals via numpy, safe to call before JIT).

Parameters:

Name	Type	Description	Default
data		A :class:SparseDiaData or dense array of shape (*batch, n, n).	required

Returns:

Type	Description
'SparseDiaImpl'	A new SparseDiaImpl instance.

Source code in jaxquantum/core/sparse_dia.py

@classmethod
def from_data(cls, data) -> "SparseDiaImpl":
    """Wrap *data* in a new ``SparseDiaImpl``.

    Accepts either a :class:`SparseDiaData` container (direct wrap) or
    a dense array-like (auto-detect non-zero diagonals via numpy, safe
    to call before JIT).

    Args:
        data: A :class:`SparseDiaData` or dense array of shape
            (*batch, n, n).

    Returns:
        A new ``SparseDiaImpl`` instance.
    """
    if isinstance(data, SparseDiaData):
        return cls._make(data.offsets, data.diags)
    offsets, diags_np = _dense_to_sparsedia(np.asarray(data))
    return cls._make(offsets, jnp.array(diags_np))

from_diags(offsets, diags) classmethod

Directly construct from sorted offsets and padded diagonal array.

This is the preferred factory when diagonal structure is known in advance (e.g., when building destroy or create operators).

Parameters:

Name	Type	Description	Default
offsets	tuple	Tuple of integer diagonal offsets (need not be sorted; will be sorted internally).	required
diags	Array	JAX array of shape (batch, n_diags, n) with padded diagonal values matching *offsets.	required

Returns:

Type	Description
'SparseDiaImpl'	A new SparseDiaImpl instance.

Source code in jaxquantum/core/sparse_dia.py

@classmethod
def from_diags(cls, offsets: tuple, diags: Array) -> "SparseDiaImpl":
    """Directly construct from sorted offsets and padded diagonal array.

    This is the preferred factory when diagonal structure is known in
    advance (e.g., when building ``destroy`` or ``create`` operators).

    Args:
        offsets: Tuple of integer diagonal offsets (need not be sorted;
            will be sorted internally).
        diags:   JAX array of shape (*batch, n_diags, n) with padded
                 diagonal values matching *offsets*.

    Returns:
        A new ``SparseDiaImpl`` instance.
    """
    return cls._make(tuple(sorted(offsets)), diags)

get_data()

Return a :class:SparseDiaData container with the raw diagonal data.

Source code in jaxquantum/core/sparse_dia.py

def get_data(self) -> SparseDiaData:
    """Return a :class:`SparseDiaData` container with the raw diagonal data."""
    return SparseDiaData(offsets=self._offsets, diags=self._diags)

imag()

Element-wise imaginary part of stored values.

Source code in jaxquantum/core/sparse_dia.py

def imag(self) -> "SparseDiaImpl":
    """Element-wise imaginary part of stored values."""
    return SparseDiaImpl._make(
        self._offsets,
        jnp.imag(self._diags).astype(self._diags.dtype),
    )

kron(other)

Kronecker product.

SparseDIA ⊗ SparseDIA stays SparseDIA: output offset for pair (kA, kB) is kA * m + kB where m = dim(B). Fully vectorised — no loops at JAX level.

Source code in jaxquantum/core/sparse_dia.py

def kron(self, other: QarrayImpl) -> QarrayImpl:
    """Kronecker product.

    SparseDIA ⊗ SparseDIA stays SparseDIA: output offset for pair
    (kA, kB) is ``kA * m + kB`` where m = dim(B).  Fully vectorised —
    no loops at JAX level.
    """
    if isinstance(other, SparseDiaImpl):
        return _sparsedia_kron(self, other)
    a, b = self._coerce(other)
    if a is not self:
        return a.kron(b)
    return a.kron(b)

matmul(other)

Matrix multiplication.

• SparseDIA @ SparseDIA → SparseDIA (O(d₁·d₂·n))
• SparseDIA @ Dense → Dense (O(d·n²), no densification of self)
• Others → coerce then delegate

Source code in jaxquantum/core/sparse_dia.py

def matmul(self, other: QarrayImpl) -> QarrayImpl:
    """Matrix multiplication.

    * SparseDIA @ SparseDIA → SparseDIA  (O(d₁·d₂·n))
    * SparseDIA @ Dense    → Dense       (O(d·n²), no densification of self)
    * Others               → coerce then delegate
    """
    if isinstance(other, DenseImpl):
        return DenseImpl._make(_sparsedia_matmul_dense(
            self._offsets, self._diags, other._data
        ))
    if isinstance(other, SparseDiaImpl):
        offsets, diags = _sparsedia_matmul_sparsedia(
            self._offsets, self._diags,
            other._offsets, other._diags,
        )
        return SparseDiaImpl._make(offsets, diags)
    a, b = self._coerce(other)
    if a is not self:
        return a.matmul(b)
    return a.matmul(b)

mul(scalar)

Scalar multiplication — scales all diagonal values.

Source code in jaxquantum/core/sparse_dia.py

def mul(self, scalar) -> "SparseDiaImpl":
    """Scalar multiplication — scales all diagonal values."""
    return SparseDiaImpl._make(self._offsets, scalar * self._diags)

neg()

Negation.

Source code in jaxquantum/core/sparse_dia.py

def neg(self) -> "SparseDiaImpl":
    """Negation."""
    return SparseDiaImpl._make(self._offsets, -self._diags)

powm(n)

Integer matrix power staying SparseDIA via binary exponentiation.

Uses O(log n) SparseDIA @ SparseDIA multiplications rather than densifying. A^0 returns the identity operator.

Parameters:

Name	Type	Description	Default
n	int	Non-negative integer exponent.	required

Returns:

Type	Description
'SparseDiaImpl'	A SparseDiaImpl equal to this matrix raised to the n-th power.

Raises:

Type	Description
ValueError	If n is negative.

Source code in jaxquantum/core/sparse_dia.py

def powm(self, n: int) -> "SparseDiaImpl":
    """Integer matrix power staying SparseDIA via binary exponentiation.

    Uses O(log n) SparseDIA @ SparseDIA multiplications rather than
    densifying.  A^0 returns the identity operator.

    Args:
        n: Non-negative integer exponent.

    Returns:
        A ``SparseDiaImpl`` equal to this matrix raised to the *n*-th power.

    Raises:
        ValueError: If *n* is negative.
    """
    if n < 0:
        raise ValueError("powm requires n >= 0")
    if n == 0:
        size = self._diags.shape[-1]
        eye_diags = jnp.ones((*self._diags.shape[:-2], 1, size), dtype=self._diags.dtype)
        return SparseDiaImpl._make((0,), eye_diags)
    if n == 1:
        return self
    half = self.powm(n // 2)
    squared = half.matmul(half)  # SparseDIA @ SparseDIA → SparseDIA
    return squared if n % 2 == 0 else self.matmul(squared)

real()

Element-wise real part of stored values.

Source code in jaxquantum/core/sparse_dia.py

def real(self) -> "SparseDiaImpl":
    """Element-wise real part of stored values."""
    return SparseDiaImpl._make(
        self._offsets,
        jnp.real(self._diags).astype(self._diags.dtype),
    )

shape()

Shape of the represented square matrix (including batch dims).

Source code in jaxquantum/core/sparse_dia.py

def shape(self) -> tuple:
    """Shape of the represented square matrix (including batch dims)."""
    n = self._diags.shape[-1]
    return (*self._diags.shape[:-2], n, n)

sub(other)

Element-wise subtraction.

Source code in jaxquantum/core/sparse_dia.py

def sub(self, other: QarrayImpl) -> QarrayImpl:
    """Element-wise subtraction."""
    if isinstance(other, SparseDiaImpl):
        return _sparsedia_add(self, other, subtract=True)
    a, b = self._coerce(other)
    if a is not self:
        return a.sub(b)
    return a.sub(b)

tidy_up(atol)

Zero diagonal values whose magnitude is below atol.

Source code in jaxquantum/core/sparse_dia.py

def tidy_up(self, atol) -> "SparseDiaImpl":
    """Zero diagonal values whose magnitude is below *atol*."""
    diags = self._diags
    real_part = jnp.where(jnp.abs(jnp.real(diags)) < atol, 0.0, jnp.real(diags))
    if jnp.issubdtype(diags.dtype, jnp.complexfloating):
        imag_part = jnp.where(jnp.abs(jnp.imag(diags)) < atol, 0.0, jnp.imag(diags))
        new_diags = (real_part + 1j * imag_part).astype(diags.dtype)
    else:
        new_diags = real_part.astype(diags.dtype)
    return SparseDiaImpl._make(self._offsets, new_diags)

to_dense()

Convert to a DenseImpl by summing diagonal contributions.

Source code in jaxquantum/core/sparse_dia.py

def to_dense(self) -> "DenseImpl":
    """Convert to a ``DenseImpl`` by summing diagonal contributions."""
    n = self._diags.shape[-1]
    batch_shape = self._diags.shape[:-2]
    result = jnp.zeros((*batch_shape, n, n), dtype=self._diags.dtype)
    for i, k in enumerate(self._offsets):
        s = _dia_slice(k)
        length = n - abs(k)
        if length <= 0:
            continue
        vals = self._diags[..., i, s]
        row_idx = jnp.arange(length) + max(-k, 0)
        col_idx = row_idx + k
        result = result.at[..., row_idx, col_idx].set(vals)
    return DenseImpl._make(result)

to_sparse_bcoo()

Convert to a SparseBCOOImpl (BCOO) via dense.

Source code in jaxquantum/core/sparse_dia.py

def to_sparse_bcoo(self) -> "SparseBCOOImpl":
    """Convert to a ``SparseBCOOImpl`` (BCOO) via dense."""
    return self.to_dense().to_sparse_bcoo()

to_sparse_dia()

Return self (already SparseDIA).

Source code in jaxquantum/core/sparse_dia.py

def to_sparse_dia(self) -> "SparseDiaImpl":
    """Return self (already SparseDIA)."""
    return self

trace()

Compute trace directly from the main diagonal (offset 0).

Returns:

Type	Description
	Scalar trace (sum of main diagonal values).

Source code in jaxquantum/core/sparse_dia.py

def trace(self):
    """Compute trace directly from the main diagonal (offset 0).

    Returns:
        Scalar trace (sum of main diagonal values).
    """
    if 0 in self._offsets:
        i = self._offsets.index(0)
        return jnp.sum(self._diags[..., i, :], axis=-1)
    return jnp.zeros(self._diags.shape[:-2], dtype=self._diags.dtype)

basis(N, k, implementation=QarrayImplType.DENSE)

Creates a |k> (i.e. fock state) ket in a specified Hilbert Space.

Parameters:

Name	Type	Description	Default
N	int	Hilbert space dimension	required
k	int	fock number	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
	Fock State |k>

Source code in jaxquantum/core/operators.py

def basis(N: int, k: int, implementation: QarrayImplType = QarrayImplType.DENSE):
    """Creates a |k> (i.e. fock state) ket in a specified Hilbert Space.

    Args:
        N: Hilbert space dimension
        k: fock number
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        Fock State |k>
    """
    return Qarray.create(one_hot(k, N), qtype="ket", implementation=implementation)

basis_like(A, ks)

Creates a |k> (i.e. fock state) ket with the same space dims as A.

Parameters:

Name	Type	Description	Default
A	Qarray	state or operator.	required
k		fock number.	required

Returns:

Type	Description
Qarray	Fock State |k> with the same space dims as A.

Source code in jaxquantum/core/operators.py

def basis_like(A: Qarray, ks: List[int]) -> Qarray:
    """Creates a |k> (i.e. fock state) ket with the same space dims as A.

    Args:
        A: state or operator.
        k: fock number.

    Returns:
        Fock State |k> with the same space dims as A.
    """
    space_dims = A.space_dims
    assert len(space_dims) == len(ks), "len(ks) must be equal to len(space_dims)"

    kets = []
    for j, k in enumerate(ks):
        kets.append(basis(space_dims[j], k))
    return tensor(*kets)

cf_wigner(psi, xvec, yvec)

Wigner function for a state vector or density matrix at points xvec + i * yvec.

Parameters

Qarray

A state vector or density matrix.

array_like

x-coordinates at which to calculate the Wigner function.

array_like

y-coordinates at which to calculate the Wigner function.

Returns

array

Values representing the Wigner function calculated over the specified range [xvec,yvec].

Source code in jaxquantum/core/cfunctions.py

def cf_wigner(psi, xvec, yvec):
    """Wigner function for a state vector or density matrix at points
    `xvec + i * yvec`.

    Parameters
    ----------

    state : Qarray
        A state vector or density matrix.

    xvec : array_like
        x-coordinates at which to calculate the Wigner function.

    yvec : array_like
        y-coordinates at which to calculate the Wigner function.


    Returns
    -------

    W : array
        Values representing the Wigner function calculated over the specified
        range [xvec,yvec].


    """
    N = psi.dims[0][0]
    x, y = jnp.meshgrid(xvec, yvec)
    alpha = x + 1.0j * y
    displacement = jqt.displace(N, alpha)

    vmapped_overlap = [vmap(vmap(jqt.overlap, in_axes=(None, 0)), in_axes=(
        None, 0))]
    for _ in psi.bdims:
        vmapped_overlap.append(vmap(vmapped_overlap[-1], in_axes=(0, None)))

    cf = vmapped_overlap[-1](psi, displacement)
    return cf

coherent(N, α)

Coherent state.

Parameters:

Name	Type	Description	Default
N	int	Hilbert Space Size.	required
α	complex	coherent state amplitude.	required

Return

Coherent state |α⟩.

Source code in jaxquantum/core/operators.py

def coherent(N: int, α: complex) -> Qarray:
    """Coherent state.

    Args:
        N: Hilbert Space Size.
        α: coherent state amplitude.

    Return:
        Coherent state |α⟩.
    """
    return displace(N, α) @ basis(N, 0)

collapse(qarr, mode='sum')

Collapse the batch dimensions of qarr.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Quantum array with optional batch dimensions.	required
mode		Collapse strategy. Only "sum" is currently supported.	'sum'

Returns:

Type	Description
Qarray	A non-batched Qarray obtained by summing over all batch axes.

Source code in jaxquantum/core/qarray.py

def collapse(qarr: Qarray, mode="sum") -> Qarray:
    """Collapse the batch dimensions of *qarr*.

    Args:
        qarr: Quantum array with optional batch dimensions.
        mode: Collapse strategy.  Only ``"sum"`` is currently supported.

    Returns:
        A non-batched ``Qarray`` obtained by summing over all batch axes.
    """

    if mode == "sum":
        if len(qarr.bdims) == 0:
            return qarr

        batch_axes = list(range(len(qarr.bdims)))

        # Preserve implementation type
        implementation = qarr.impl_type
        return Qarray.create(jnp.sum(qarr.data, axis=batch_axes), dims=qarr.dims, implementation=implementation)

concatenate(qarr_list, axis=0)

Concatenate a list of Qarrays along a specified axis.

Parameters:

Name	Type	Description	Default
qarr_list	List[Qarray]	List of Qarrays to concatenate.	required
axis	int	Axis along which to concatenate. Default is 0.	0

Returns:

Type	Description
Qarray	Concatenated Qarray.

Source code in jaxquantum/core/qarray.py

def concatenate(qarr_list: List[Qarray], axis: int = 0) -> Qarray:
    """Concatenate a list of Qarrays along a specified axis.

    Args:
        qarr_list: List of Qarrays to concatenate.
        axis: Axis along which to concatenate. Default is 0.

    Returns:
        Concatenated Qarray.
    """

    non_empty_qarr_list = [qarr for qarr in qarr_list if len(qarr.data) != 0]

    if len(non_empty_qarr_list) == 0:
        return Qarray.from_list([])

    concatenated_data = jnp.concatenate(
        [qarr.data for qarr in non_empty_qarr_list], axis=axis
    )

    dims = non_empty_qarr_list[0].dims
    return Qarray.create(concatenated_data, dims=dims)

cosm(qarr)

Matrix cosine of a Qarray.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array (converted to dense internally).	required

Returns:

Type	Description
Qarray	A dense Qarray containing the matrix cosine.

Source code in jaxquantum/core/qarray.py

def cosm(qarr: Qarray) -> Qarray:
    """Matrix cosine of a ``Qarray``.

    Args:
        qarr: Input quantum array (converted to dense internally).

    Returns:
        A dense ``Qarray`` containing the matrix cosine.
    """
    dims = qarr.dims
    # Convert to dense for cosm
    dense_data = qarr.to_dense().data
    data = cosm_data(dense_data)
    return Qarray.create(data, dims=dims)

cosm_data(data, **kwargs)

Matrix cosine of a raw array.

Parameters:

Name	Type	Description	Default
data	Array	Dense matrix array.	required
**kwargs		Unused; kept for API consistency.	{}

Returns:

Type	Description
Array	The matrix cosine computed as (expm(i*A) + expm(-i*A)) / 2.

Source code in jaxquantum/core/qarray.py

def cosm_data(data: Array, **kwargs) -> Array:
    """Matrix cosine of a raw array.

    Args:
        data: Dense matrix array.
        **kwargs: Unused; kept for API consistency.

    Returns:
        The matrix cosine computed as ``(expm(i*A) + expm(-i*A)) / 2``.
    """
    return (expm_data(1j * data) + expm_data(-1j * data)) / 2

create(N, implementation=QarrayImplType.DENSE)

creation operator

Parameters:

Name	Type	Description	Default
N		Hilbert space size	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	creation operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def create(N, implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """creation operator

    Args:
        N: Hilbert space size
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        creation operator in Hilber Space of size N
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # Single subdiagonal at offset -1; Convention A: 1 trailing zero.
        diags = jnp.zeros((1, N), dtype=jnp.float64)
        diags = diags.at[0, :N - 1].set(jnp.sqrt(jnp.arange(1, N, dtype=jnp.float64)))
        return _make_sparsedia(offsets=(-1,), diags=diags)
    return Qarray.create(jnp.diag(jnp.sqrt(jnp.arange(1, N)), k=-1), implementation=implementation)

dag(qarr)

Conjugate transpose of qarr.

For ket/bra vectors (stored on a single axis) this is just a complex conjugate with the dims reversed — no axis swap, so the data stays 1‑D. For operators it is the usual conjugate transpose.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array.	required

Returns:

Type	Description
Qarray	The conjugate transpose with swapped dims.

Source code in jaxquantum/core/qarray.py

def dag(qarr: Qarray) -> Qarray:
    """Conjugate transpose of *qarr*.

    For ket/bra vectors (stored on a single axis) this is just a complex
    conjugate with the ``dims`` reversed — no axis swap, so the data stays
    1‑D.  For operators it is the usual conjugate transpose.

    Args:
        qarr: Input quantum array.

    Returns:
        The conjugate transpose with swapped ``dims``.
    """
    dims = qarr.dims[::-1]
    if qarr.qtype in (Qtypes.ket, Qtypes.bra):
        new_impl = qarr._impl.conj()
    else:
        new_impl = qarr._impl.dag()
    # Infer batch dimensions from the transformed data because ``_bdims`` is
    # static PyTree metadata and can be stale under ``vmap``.
    return Qarray._from_impl(new_impl, Qdims(dims))

dag_data(arr)

Conjugate transpose of a raw array, dispatching to the right backend.

Iterates through registered :class:QarrayImpl subclasses and delegates to the first one whose :meth:~QarrayImpl.can_handle_data returns True. Adding a new backend automatically extends this function — no changes required here.

Parameters:

Name	Type	Description	Default
arr		Input array (jnp.ndarray, sparse.BCOO, or any type handled by a registered impl). For 1-D dense arrays only conjugation is applied (no transpose).	required

Returns:

Type	Description
Array	Conjugate transpose with the last two axes swapped.

Raises:

Type	Description
TypeError	If no registered impl can handle arr.

Source code in jaxquantum/core/qarray.py

def dag_data(arr) -> Array:
    """Conjugate transpose of a raw array, dispatching to the right backend.

    Iterates through registered :class:`QarrayImpl` subclasses and delegates
    to the first one whose :meth:`~QarrayImpl.can_handle_data` returns True.
    Adding a new backend automatically extends this function — no changes
    required here.

    Args:
        arr: Input array (``jnp.ndarray``, ``sparse.BCOO``, or any type
            handled by a registered impl).  For 1-D dense arrays only
            conjugation is applied (no transpose).

    Returns:
        Conjugate transpose with the last two axes swapped.

    Raises:
        TypeError: If no registered impl can handle *arr*.
    """
    for impl_class in _IMPL_REGISTRY:
        if impl_class.can_handle_data(arr):
            return impl_class.dag_data(arr)
    raise TypeError(f"dag_data: no registered impl can handle type {type(arr)}")

destroy(N, implementation=QarrayImplType.DENSE)

annihilation operator

Parameters:

Name	Type	Description	Default
N		Hilbert space size	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	annilation operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def destroy(N, implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """annihilation operator

    Args:
        N: Hilbert space size
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        annilation operator in Hilber Space of size N
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # Single superdiagonal at offset +1; Convention A: 1 leading zero.
        diags = jnp.zeros((1, N), dtype=jnp.float64)
        diags = diags.at[0, 1:].set(jnp.sqrt(jnp.arange(1, N, dtype=jnp.float64)))
        return _make_sparsedia(offsets=(1,), diags=diags)
    return Qarray.create(jnp.diag(jnp.sqrt(jnp.arange(1, N)), k=1), implementation=implementation)

displace(N, α)

Displacement operator

Parameters:

Name	Type	Description	Default
N		Hilbert Space Size	required
α		Phase space displacement	required

Returns:

Type	Description
Qarray	Displace operator D(α)

Source code in jaxquantum/core/operators.py

def displace(N, α) -> Qarray:
    """Displacement operator

    Args:
        N: Hilbert Space Size
        α: Phase space displacement

    Returns:
        Displace operator D(α)
    """
    a = destroy(N)
    return (α * a.dag() - jnp.conj(α) * a).expm()

eigenenergies(qarr)

Eigenvalues of a quantum array.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Hermitian operator (converted to dense internally).	required

Returns:

Type	Description
Array	Sorted eigenvalues as a JAX array.

Source code in jaxquantum/core/qarray.py

def eigenenergies(qarr: Qarray) -> Array:
    """Eigenvalues of a quantum array.

    Args:
        qarr: Hermitian operator (converted to dense internally).

    Returns:
        Sorted eigenvalues as a JAX array.
    """
    # Convert to dense for eigenenergies
    dense_qarr = qarr.to_dense()
    evals = jnp.linalg.eigvalsh(dense_qarr.data)
    return evals

eigenstates(qarr)

Eigenstates of a quantum array.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Hermitian operator (converted to dense internally).	required

Returns:

Type	Description
Qarray	A tuple (eigenvalues, eigenstates_qarray) where eigenvalues are
Qarray	sorted in ascending order.

Source code in jaxquantum/core/qarray.py

def eigenstates(qarr: Qarray) -> Qarray:
    """Eigenstates of a quantum array.

    Args:
        qarr: Hermitian operator (converted to dense internally).

    Returns:
        A tuple ``(eigenvalues, eigenstates_qarray)`` where eigenvalues are
        sorted in ascending order.
    """
    # Convert to dense for eigenstates
    dense_qarr = qarr.to_dense()

    evals, evecs = jnp.linalg.eigh(dense_qarr.data)
    dims = ket_from_op_dims(qarr.dims)

    # numpy returns [batch, :, i] as the i-th eigenvector
    # we want [batch, i, :] as the i-th eigenvector
    evecs = jnp.swapaxes(evecs, -2, -1)
    evecs = Qarray._from_impl(
        DenseImpl._make(evecs),
        Qdims(dims),
    )

    return evals, evecs

expm(qarr, **kwargs)

Matrix exponential of a Qarray.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array (converted to dense internally).	required
**kwargs		Forwarded to jsp.linalg.expm.	{}

Returns:

Type	Description
Qarray	A dense Qarray containing the matrix exponential.

Source code in jaxquantum/core/qarray.py

def expm(qarr: Qarray, **kwargs) -> Qarray:
    """Matrix exponential of a ``Qarray``.

    Args:
        qarr: Input quantum array (converted to dense internally).
        **kwargs: Forwarded to ``jsp.linalg.expm``.

    Returns:
        A dense ``Qarray`` containing the matrix exponential.
    """
    dims = qarr.dims
    # Convert to dense for expm
    dense_data = qarr.to_dense().data
    data = expm_data(dense_data, **kwargs)
    return Qarray.create(data, dims=dims)

expm_data(data, **kwargs)

Matrix exponential of a raw array.

Parameters:

Name	Type	Description	Default
data	Array	Dense matrix array.	required
**kwargs		Forwarded to jsp.linalg.expm.	{}

Returns:

Type	Description
Array	The matrix exponential.

Source code in jaxquantum/core/qarray.py

def expm_data(data: Array, **kwargs) -> Array:
    """Matrix exponential of a raw array.

    Args:
        data: Dense matrix array.
        **kwargs: Forwarded to ``jsp.linalg.expm``.

    Returns:
        The matrix exponential.
    """
    return jsp.linalg.expm(data, **kwargs)

extract_dims(arr, dims=None)

Extract dims from a JAX array or Qarray.

Parameters:

Name	Type	Description	Default
arr	Array	JAX array or Qarray.	required
dims	Optional[Union[DIMS_TYPE, List[int]]]	Qarray dims.	None

Returns:

Type	Description
	Qarray dims.

Source code in jaxquantum/core/conversions.py

def extract_dims(arr: Array, dims: Optional[Union[DIMS_TYPE, List[int]]] = None):
    """Extract dims from a JAX array or Qarray.

    Args:
        arr: JAX array or Qarray.
        dims: Qarray dims.

    Returns:
        Qarray dims.
    """
    if isinstance(dims[0], Number):
        # A 1-D array is a ket; a square 2-D array is an operator; anything else
        # (non-square 2-D) is a batch of kets. Vectors are stored as (N,).
        is_op = arr.ndim >= 2 and arr.shape[-2] == arr.shape[-1]
        if is_op:
            dims = [dims, dims]
        else:
            dims = [dims, [1] * len(dims)]  # defaults to ket
    return dims

fidelity(rho, sigma, force_positivity=False)

Fidelity between two states.

Parameters:

Name	Type	Description	Default
rho	Qarray	state.	required
sigma	Qarray	state.	required
force_positivity	bool	force the states to be positive semidefinite	False

Returns:

Type	Description
ndarray	Fidelity between rho and sigma.

Source code in jaxquantum/core/measurements.py

def fidelity(rho: Qarray, sigma: Qarray, force_positivity: bool=False) -> (
        jnp.ndarray):
    """Fidelity between two states.

    Args:
        rho: state.
        sigma: state.
        force_positivity: force the states to be positive semidefinite

    Returns:
        Fidelity between rho and sigma.
    """
    rho = rho.to_dm()
    sigma = sigma.to_dm()

    sqrt_rho = powm(rho, 0.5, clip_eigvals=force_positivity)

    return jnp.real(((powm(sqrt_rho @ sigma @ sqrt_rho, 0.5,
                           clip_eigvals=force_positivity)).tr())
                    ** 2)

hadamard(implementation=QarrayImplType.DENSE)

H

Returns:

Name	Type	Description
H	Qarray	Hadamard gate

Source code in jaxquantum/core/operators.py

def hadamard(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """H

    Returns:
        H: Hadamard gate
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        s = 1.0 / jnp.sqrt(2.0)
        # offset -1: valid at [0]   → diag[0]=A[1,0]=s, diag[1]=0 (trailing zero)
        # offset  0: valid at [0:2] → diag[0]=A[0,0]=s, diag[1]=A[1,1]=-s
        # offset +1: valid at [1]   → diag[0]=0 (leading zero), diag[1]=A[0,1]=s
        diags = jnp.array([[s, 0.0], [s, -s], [0.0, s]])
        return _make_sparsedia(offsets=(-1, 0, 1), diags=diags)
    return Qarray.create(jnp.array([[1, 1], [1, -1]]) / jnp.sqrt(2), implementation=implementation)

identity(*args, implementation=QarrayImplType.DENSE, **kwargs)

Identity matrix.

Parameters:

Name	Type	Description	Default
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	Identity matrix.

Source code in jaxquantum/core/operators.py

def identity(*args, implementation: QarrayImplType = QarrayImplType.DENSE, **kwargs) -> Qarray:
    """Identity matrix.

    Args:
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        Identity matrix.
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # jnp.eye(*args) is typically eye(N) or eye(N, N); extract N from args.
        n = args[0] if args else kwargs.get("N", kwargs.get("n", None))
        if n is not None and (len(args) <= 1) and not kwargs:
            diags = jnp.ones((1, int(n)), dtype=jnp.float64)
            return _make_sparsedia(offsets=(0,), diags=diags)
    return Qarray.create(jnp.eye(*args, **kwargs), implementation=implementation)

identity_like(A, implementation=QarrayImplType.DENSE)

Identity matrix with the same shape as A.

Parameters:

Name	Type	Description	Default
A		Matrix.	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	Identity matrix with the same shape as A.

Source code in jaxquantum/core/operators.py

def identity_like(A, implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """Identity matrix with the same shape as A.

    Args:
        A: Matrix.
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        Identity matrix with the same shape as A.
    """
    space_dims = A.space_dims
    total_dim = prod(space_dims)
    return Qarray.create(jnp.eye(total_dim, total_dim), dims=[space_dims, space_dims], implementation=implementation)

is_dm_data(data)

Check whether data has the shape of a density matrix (square matrix).

Parameters:

Name	Type	Description	Default
data	Array	Array to check.	required

Returns:

Type	Description
bool	True if the last two dimensions are equal.

Source code in jaxquantum/core/qarray.py

def is_dm_data(data: Array) -> bool:
    """Check whether *data* has the shape of a density matrix (square matrix).

    Args:
        data: Array to check.

    Returns:
        True if the last two dimensions are equal.
    """
    return data.shape[-2] == data.shape[-1]

jnp2jqt(arr, dims=None)

JAX array -> QuTiP state.

Parameters:

Name	Type	Description	Default
jnp_obj		JAX array.	required
dims	Optional[Union[DIMS_TYPE, List[int]]]	Qarray dims.	None

Returns:

Type	Description
	QuTiP state.

Source code in jaxquantum/core/conversions.py

def jnp2jqt(arr: Array, dims: Optional[Union[DIMS_TYPE, List[int]]] = None):
    """JAX array -> QuTiP state.

    Args:
        jnp_obj: JAX array.
        dims: Qarray dims.

    Returns:
        QuTiP state.
    """
    dims = extract_dims(arr, dims) if dims is not None else None
    return Qarray.create(arr, dims=dims)

jqt2qt(jqt_obj)

Qarray -> QuTiP state.

Parameters:

Name	Type	Description	Default
jqt_obj		Qarray.	required
dims		QuTiP dims.	required

Returns:

Type	Description
	QuTiP state.

Source code in jaxquantum/core/conversions.py

def jqt2qt(jqt_obj):
    """Qarray -> QuTiP state.

    Args:
        jqt_obj: Qarray.
        dims: QuTiP dims.

    Returns:
        QuTiP state.
    """
    if isinstance(jqt_obj, Qobj) or jqt_obj is None:
        return jqt_obj

    if jqt_obj.is_batched:
        res = []
        for i in range(len(jqt_obj)):
            res.append(jqt2qt(jqt_obj[i]))
        return res

    dims = [list(jqt_obj.dims[0]), list(jqt_obj.dims[1])]
    return Qobj(np.array(jqt_obj.data), dims=dims)

keep_only_diag_elements(qarr)

Zero out all off-diagonal elements of qarr.

For sparse Qarray objects the off-diagonal stored values are zeroed in-place on the BCOO structure — no densification.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Non-batched input quantum array.	required

Returns:

Type	Description
Qarray	A Qarray with only diagonal entries non-zero.

Raises:

Type	Description
ValueError	If qarr has batch dimensions.

Source code in jaxquantum/core/qarray.py

def keep_only_diag_elements(qarr: Qarray) -> Qarray:
    """Zero out all off-diagonal elements of *qarr*.

    For sparse ``Qarray`` objects the off-diagonal stored values are zeroed
    in-place on the BCOO structure — no densification.

    Args:
        qarr: Non-batched input quantum array.

    Returns:
        A ``Qarray`` with only diagonal entries non-zero.

    Raises:
        ValueError: If *qarr* has batch dimensions.
    """
    if len(qarr.bdims) > 0:
        raise ValueError("Cannot keep only diagonal elements of a batched Qarray.")

    dims = qarr.dims
    if qarr.is_sparse_bcoo:
        new_impl = qarr._impl.keep_only_diag()
        return Qarray.create(new_impl.data, dims=dims, implementation=QarrayImplType.SPARSE_BCOO)
    if qarr.is_sparse_dia:
        from jaxquantum.core.sparse_dia import SparseDiaImpl
        impl = qarr._impl
        n = impl._diags.shape[-1]
        if 0 in impl._offsets:
            i = impl._offsets.index(0)
            main_diag = impl._diags[..., i:i + 1, :]
        else:
            main_diag = jnp.zeros((*impl._diags.shape[:-2], 1, n), dtype=impl._diags.dtype)
        new_impl = SparseDiaImpl(_offsets=(0,), _diags=main_diag)
        return Qarray.create(new_impl.get_data(), dims=dims, implementation=QarrayImplType.SPARSE_DIA)
    data = jnp.diag(jnp.diag(qarr.data))
    return Qarray.create(data, dims=dims)

ket2dm(qarr)

Convert a ket to a density matrix via outer product.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Ket, bra, or operator. Operators are returned unchanged.	required

Returns:

Type	Description
Qarray	Density matrix |ψ⟩⟨ψ|.

Source code in jaxquantum/core/qarray.py

def ket2dm(qarr: Qarray) -> Qarray:
    """Convert a ket to a density matrix via outer product.

    Args:
        qarr: Ket, bra, or operator.  Operators are returned unchanged.

    Returns:
        Density matrix ``|ψ⟩⟨ψ|``.
    """
    if qarr.qtype == Qtypes.oper:
        return qarr

    if qarr.qtype == Qtypes.bra:
        qarr = qarr.dag()

    return qarr @ qarr.dag()

mesolve(H, rho0, tlist, saveat_tlist=None, c_ops=None, solver_options=None)

Solve a Lindblad master equation and return the saved states.

Parameters:

Name	Type	Description	Default
H	Qarray | Callable[[float], Qarray]	Static Hamiltonian or callable H(t).	required
rho0	Qarray	Initial ket or density matrix.	required
tlist	Array	Integration interval; also the default save times.	required
saveat_tlist	Array | None	Save times. An empty array saves only the final state.	None
c_ops	Qarray | None	Collapse operators.	None
solver_options	SolverOptions | None	Native Diffrax configuration.	None

Returns:

Type	Description
Qarray	Saved density matrices as a batched Qarray.

See Also

:func:mesolve_result returns the complete Diffrax solution.

Source code in jaxquantum/core/solvers.py

def mesolve(
    H: Qarray | Callable[[float], Qarray],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Array | None = None,
    c_ops: Qarray | None = None,
    solver_options: SolverOptions | None = None,
) -> Qarray:
    """Solve a Lindblad master equation and return the saved states.

    Args:
        H: Static Hamiltonian or callable ``H(t)``.
        rho0: Initial ket or density matrix.
        tlist: Integration interval; also the default save times.
        saveat_tlist: Save times. An empty array saves only the final state.
        c_ops: Collapse operators.
        solver_options: Native Diffrax configuration.

    Returns:
        Saved density matrices as a batched ``Qarray``.

    See Also:
        :func:`mesolve_result` returns the complete Diffrax solution.
    """
    solution = mesolve_result(
        H,
        rho0,
        tlist,
        saveat_tlist=saveat_tlist,
        c_ops=c_ops,
        solver_options=solver_options,
    )
    qdims = Qdims((rho0.space_dims, rho0.space_dims))
    return Qarray._from_impl(DenseImpl._make(solution.ys), qdims)

mesolve_result(H, rho0, tlist, saveat_tlist=None, c_ops=None, solver_options=None)

Solve a Lindblad master equation and return its Diffrax solution.

Use this form for solver statistics, events, dense interpolation, custom SaveAt functions, or continuation state.

Source code in jaxquantum/core/solvers.py

def mesolve_result(
    H: Qarray | Callable[[float], Qarray],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Array | None = None,
    c_ops: Qarray | None = None,
    solver_options: SolverOptions | None = None,
) -> diffrax.Solution:
    """Solve a Lindblad master equation and return its Diffrax solution.

    Use this form for solver statistics, events, dense interpolation, custom
    ``SaveAt`` functions, or continuation state.
    """
    collapse_ops = c_ops if c_ops is not None else Qarray.from_list([])

    if len(collapse_ops) == 0 and rho0.qtype != Qtypes.oper:
        logging.warning(  # noqa: LOG015
            "Consider sesolve(): no collapse operators were provided and the "
            "initial state is not a density matrix."
        )

    rho_data = rho0.to_dm().to_dense()

    if robust_isscalar(H):
        H = H * identity_like(rho_data)

    if isinstance(H, Qarray):
        H_data = lambda t: H.data
    else:
        H_data = lambda t: H(t).data

    return _mesolve_result_data(
        H_data,
        rho_data.data,
        tlist,
        saveat_tlist,
        collapse_ops.data,
        solver_options=solver_options,
    )

multi_mode_basis_set(Ns)

Creates a multi-mode basis set.

Parameters:

Name	Type	Description	Default
Ns	List[int]	List of Hilbert space dimensions for each mode.	required

Returns:

Type	Description
Qarray	Multi-mode basis set.

Source code in jaxquantum/core/operators.py

def multi_mode_basis_set(Ns: List[int]) -> Qarray:
    """Creates a multi-mode basis set.

    Args:
        Ns: List of Hilbert space dimensions for each mode.

    Returns:
        Multi-mode basis set.
    """
    data = jnp.eye(prod(Ns))
    dims = (tuple(Ns), tuple([1 for _ in Ns]))
    return Qarray.create(data, dims=dims, bdims=(prod(Ns),))

norm(qarr)

Compute the norm of a quantum array.

Sparse paths (no densification):

• ket / bra — L2 norm via :meth:SparseBCOOImpl.l2_norm_batched (handles batch dimensions).
• operator — trace norm assuming PSD (nuclear norm = tr(rho) for density matrices). This is exact for density matrices; for general non-PSD operators convert to dense first.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array.	required

Returns:

Type	Description
float	The norm as a scalar (or batched array of scalars).

Source code in jaxquantum/core/qarray.py

def norm(qarr: Qarray) -> float:
    """Compute the norm of a quantum array.

    Sparse paths (no densification):

    * ket / bra — L2 norm via :meth:`SparseBCOOImpl.l2_norm_batched` (handles
      batch dimensions).
    * operator — trace norm assuming PSD (nuclear norm = tr(rho) for density
      matrices).  This is exact for density matrices; for general non-PSD
      operators convert to dense first.

    Args:
        qarr: Input quantum array.

    Returns:
        The norm as a scalar (or batched array of scalars).
    """
    if qarr.qtype in [Qtypes.ket, Qtypes.bra] and qarr.is_sparse_bcoo:
        return qarr._impl.l2_norm_batched(qarr.bdims)

    if qarr.qtype == Qtypes.oper and qarr.is_sparse_bcoo:
        # Nuclear norm = trace for positive-semidefinite (density matrix) operators.
        # jnp.real strips any floating-point imaginary artefact.
        return jnp.real(qarr._impl.trace())

    if qarr.qtype == Qtypes.oper and qarr.is_sparse_dia:
        return jnp.real(qarr._impl.trace())

    qarr = qarr.to_dense()

    qdata = qarr.data

    if qarr.qtype == Qtypes.oper:
        if default_backend() == "cpu":
            return jnp.sum(jnp.linalg.svd(qdata, compute_uv=False), axis=-1)
        gram = qdata @ jnp.swapaxes(jnp.conj(qdata), -1, -2)
        values = jnp.linalg.eigvalsh(gram)
        return jnp.sum(jnp.sqrt(jnp.abs(values)), axis=-1)

    elif qarr.qtype in [Qtypes.ket, Qtypes.bra]:
        # Vectors store the Hilbert space on the single trailing axis.
        return jnp.sqrt(jnp.sum(jnp.abs(qdata) ** 2, axis=-1))

num(N, implementation=QarrayImplType.DENSE)

Number operator

Parameters:

Name	Type	Description	Default
N		Hilbert Space size	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	number operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def num(N, implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """Number operator

    Args:
        N: Hilbert Space size
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        number operator in Hilber Space of size N
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # Main diagonal only; no leading/trailing zeros needed (offset 0).
        diags = jnp.arange(N, dtype=jnp.float64).reshape(1, N)
        return _make_sparsedia(offsets=(0,), diags=diags)
    return Qarray.create(jnp.diag(jnp.arange(N)), implementation=implementation)

overlap(rho, sigma)

Overlap between two states or operators.

Parameters:

Name	Type	Description	Default
rho	Qarray	state/operator.	required
sigma	Qarray	state/operator.	required

Returns:

Type	Description
Array	Overlap between rho and sigma.

Source code in jaxquantum/core/measurements.py

def overlap(rho: Qarray, sigma: Qarray) -> Array:
    """Overlap between two states or operators.

    Args:
        rho: state/operator.
        sigma: state/operator.

    Returns:
        Overlap between rho and sigma.
    """

    if rho.is_vec() and sigma.is_vec():
        # |<a|b>|^2. Compute the inner product directly over the trailing space
        # axis (vectors are stored as (..., N)); robust for batched states.
        a = rho.to_ket().data
        b = sigma.to_ket().data
        inner = jnp.sum(jnp.conj(a) * b, axis=-1)
        return jnp.abs(inner) ** 2
    elif rho.is_vec():
        # <psi|sigma|psi>
        psi = rho.to_ket()
        Opsi = (sigma @ psi).data
        return jnp.sum(jnp.conj(psi.data) * Opsi, axis=-1)
    elif sigma.is_vec():
        # <psi|rho|psi>
        psi = sigma.to_ket()
        Opsi = (rho @ psi).data
        return jnp.sum(jnp.conj(psi.data) * Opsi, axis=-1)
    else:
        return (rho.dag() @ sigma).trace()

plot_cf(state, pts_x, pts_y=None, axs=None, contour=True, qp_type=WIGNER, cbar_label='', axis_scale_factor=1, plot_cbar=True, plot_grid=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot a characteristic function as paired real/imag subplots.

Each batch element produces two adjacent subplots — real part followed by imaginary part — so the rendered grid has shape (rows, 2 * cols).

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate the characteristic function at	required
pts_y		y points to evaluate the characteristic function at	None
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
qp_type		type of characteristic function. Currently only "wigner" is supported.	WIGNER
cbar_label		labels for the real and imaginary cbar (overridden internally based on qp_type)	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
plot_grid		whether to draw gridlines on each subplot	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots (shape must match state.bdims)	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled grid. Returns a matplotlib.animation.FuncAnimation that auto-renders inline in Jupyter.	False
gif_params		dict of options for the gif path. Recognized keys: save_path (default None) — if set, save the animation here via PillowWriter; interval_ms (default 200) — milliseconds per frame; ts (default None) — optional 1D array of timestamps matching the animation-axis length; when set, each frame's suptitle gets a t = … label; batch_animation_axis (default 0) — index into state.bdims selecting which axis becomes the animation/time axis (the remaining batch dims form the per-frame subplot grid).	None

Returns:

Type	Description
	(axs, im) in the static case, or a FuncAnimation when
	gif=True.

Source code in jaxquantum/core/visualization.py

def plot_cf(
        state,
        pts_x,
        pts_y=None,
        axs=None,
        contour=True,
        qp_type=WIGNER,
        cbar_label="",
        axis_scale_factor=1,
        plot_cbar=True,
        plot_grid=True,
        x_ticks=None,
        y_ticks=None,
        z_ticks=None,
        subtitles=None,
        figtitle=None,
        gif=False,
        gif_params=None,
):
    """Plot a characteristic function as paired real/imag subplots.

    Each batch element produces two adjacent subplots — real part followed
    by imaginary part — so the rendered grid has shape ``(rows, 2 * cols)``.

    Args:
        state: state with arbitrary number of batch dimensions, result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate the characteristic function at
        pts_y: y points to evaluate the characteristic function at
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        qp_type: type of characteristic function. Currently only
            ``"wigner"`` is supported.
        cbar_label: labels for the real and imaginary cbar (overridden
            internally based on ``qp_type``)
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        plot_grid: whether to draw gridlines on each subplot
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots (shape must match ``state.bdims``)
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled grid. Returns a ``matplotlib.animation.FuncAnimation``
            that auto-renders inline in Jupyter.
        gif_params: dict of options for the gif path. Recognized keys:
            ``save_path`` (default None) — if set, save the animation here
            via PillowWriter; ``interval_ms`` (default 200) — milliseconds
            per frame; ``ts`` (default None) — optional 1D array of
            timestamps matching the animation-axis length; when set, each
            frame's suptitle gets a ``t = …`` label;
            ``batch_animation_axis`` (default 0) — index into
            ``state.bdims`` selecting which axis becomes the animation/time
            axis (the remaining batch dims form the per-frame subplot grid).

    Returns:
        ``(axs, im)`` in the static case, or a ``FuncAnimation`` when
        ``gif=True``.
    """
    if pts_y is None:
        pts_y = pts_x
    pts_x = jnp.array(pts_x)
    pts_y = jnp.array(pts_y)

    if gif:
        return _plot_cf_gif(
            state=state,
            pts_x=pts_x,
            pts_y=pts_y,
            axs=axs,
            contour=contour,
            qp_type=qp_type,
            axis_scale_factor=axis_scale_factor,
            plot_cbar=plot_cbar,
            plot_grid=plot_grid,
            x_ticks=x_ticks,
            y_ticks=y_ticks,
            z_ticks=z_ticks,
            subtitles=subtitles,
            figtitle=figtitle,
            gif_params=gif_params or {},
        )

    bdims = state.bdims
    added_baxes = 0

    if subtitles is not None and subtitles.shape != bdims:
        raise ValueError(
            f"labels must have same shape as bdims, "
            f"got shapes {subtitles.shape} and {bdims}"
        )

    if len(bdims) == 0:
        bdims = (1,)
        added_baxes += 1
    if len(bdims) == 1:
        bdims = (1, bdims[0])
        added_baxes += 1

    extra_dims = bdims[2:]
    if extra_dims != ():
        state = state.reshape_bdims(
            bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
        )
        if subtitles is not None:
            subtitles = subtitles.reshape(
                bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
            )
        bdims = state.bdims

    if axs is None:
        _, axs = plt.subplots(
            bdims[0],
            bdims[1]*2,
            figsize=(3.3 * bdims[1]*2, 3 * bdims[0]),
            dpi=200,
        )


    if qp_type == WIGNER:
        vmin = -1
        vmax = 1
        scale = 1
        cmap = "seismic"
        cbar_label = [
            r"$\mathcal{Re}(\chi_W(\alpha))$",
            r"$\mathcal{Im}(\chi_W(\alpha))$",
        ]
        QP = scale * cf_wigner(state, pts_x, pts_y)

    for _ in range(added_baxes):
        QP = jnp.array([QP])
        axs = np.array([axs])
        if subtitles is not None:
            subtitles = np.array([subtitles])

    if added_baxes==2:
        axs = axs[0] # When the input state is zero-dimensional, remove an
                     # axis that is automatically added due to the subcolumns


    pts_x = pts_x * axis_scale_factor
    pts_y = pts_y * axis_scale_factor

    x_ticks = (
        jnp.linspace(jnp.min(pts_x), jnp.max(pts_x),
                     5) if x_ticks is None else x_ticks
    )
    y_ticks = (
        jnp.linspace(jnp.min(pts_y), jnp.max(pts_y),
                     5) if y_ticks is None else y_ticks
    )
    z_ticks = jnp.linspace(vmin, vmax, 11) if z_ticks is None else z_ticks

    im = _render_cf_grid(
        axs,
        QP,
        pts_x,
        pts_y,
        contour=contour,
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        cbar_label=cbar_label,
        plot_cbar=plot_cbar,
        plot_grid=plot_grid,
        subtitles=subtitles,
        decorate=True,
    )

    fig = axs[0, 0].get_figure()
    fig.tight_layout(w_pad=0.3, h_pad=0.3)
    if figtitle is not None:
        fig.suptitle(figtitle, y=1.04)
    return axs, im

plot_cf_wigner(state, pts_x, pts_y=None, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, plot_grid=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot the Wigner characteristic function of the state.

Thin wrapper around :func:plot_cf with qp_type='wigner'. Each batch element is rendered as two subplots side-by-side: real then imaginary part of the characteristic function.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate the characteristic function at	required
pts_y		y points to evaluate the characteristic function at	None
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
plot_grid		whether to draw gridlines on each subplot	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled subplot grid. See :func:plot_cf for details.	False
gif_params		dict of options for the gif path. Recognized keys: save_path (default None), interval_ms (default 200), ts (default None — adds a t = … label per frame), batch_animation_axis (default 0).	None

Returns:

Type	Description
	(axs, im) in the static case, or a matplotlib.animation.FuncAnimation
	when gif=True.

Source code in jaxquantum/core/visualization.py

def plot_cf_wigner(
    state,
    pts_x,
    pts_y=None,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    plot_grid=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
    gif=False,
    gif_params=None,
):
    """Plot the Wigner characteristic function of the state.

    Thin wrapper around :func:`plot_cf` with ``qp_type='wigner'``. Each batch
    element is rendered as two subplots side-by-side: real then imaginary
    part of the characteristic function.

    Args:
        state: state with arbitrary number of batch dimensions, result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate the characteristic function at
        pts_y: y points to evaluate the characteristic function at
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        plot_grid: whether to draw gridlines on each subplot
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled subplot grid. See :func:`plot_cf` for details.
        gif_params: dict of options for the gif path. Recognized keys:
            ``save_path`` (default None), ``interval_ms`` (default 200),
            ``ts`` (default None — adds a ``t = …`` label per frame),
            ``batch_animation_axis`` (default 0).

    Returns:
        ``(axs, im)`` in the static case, or a ``matplotlib.animation.FuncAnimation``
        when ``gif=True``.
    """
    return plot_cf(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        axs=axs,
        contour=contour,
        qp_type=WIGNER,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        plot_grid=plot_grid,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
        gif=gif,
        gif_params=gif_params,
    )

plot_qfunc(state, pts_x, pts_y=None, g=2, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot the husimi (Q) function of the state.

Thin wrapper around :func:plot_qp with qp_type='husimi'.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
g		float, default 2. Scaling factor for a = 0.5 * g * (x + iy). The value of g is related to the value of :math:\hbar in the commutation relation :math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	2
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled subplot grid. See :func:plot_qp for details.	False
gif_params		dict of options for the gif path. Recognized keys: save_path (default None), interval_ms (default 200), ts (default None — adds a t = … label per frame), batch_animation_axis (default 0).	None

Returns:

Type	Description
	(axs, im) in the static case, or a matplotlib.animation.FuncAnimation
	when gif=True.

Source code in jaxquantum/core/visualization.py

def plot_qfunc(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
    gif=False,
    gif_params=None,
):
    """Plot the husimi (Q) function of the state.

    Thin wrapper around :func:`plot_qp` with ``qp_type='husimi'``.

    Args:
        state: state with arbitrary number of batch dimensions, result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``.
            The value of ``g`` is related to the value of :math:`\\hbar` in
            the commutation relation :math:`[x,\,y] = i\\hbar` via
            :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled subplot grid. See :func:`plot_qp` for details.
        gif_params: dict of options for the gif path. Recognized keys:
            ``save_path`` (default None), ``interval_ms`` (default 200),
            ``ts`` (default None — adds a ``t = …`` label per frame),
            ``batch_animation_axis`` (default 0).

    Returns:
        ``(axs, im)`` in the static case, or a ``matplotlib.animation.FuncAnimation``
        when ``gif=True``.
    """
    return plot_qp(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        g=g,
        axs=axs,
        contour=contour,
        qp_type=HUSIMI,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
        gif=gif,
        gif_params=gif_params,
    )

plot_qp(state, pts_x, pts_y=None, g=2, axs=None, contour=True, qp_type=WIGNER, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot a quasi-probability distribution (Wigner or Husimi-Q).

The state may carry an arbitrary number of batch dimensions; they are flattened to a 2D (rows, cols) grid of subplots. With gif=True, one batch axis is animated instead and the remaining batch dims form the per-frame subplot grid.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions; result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate the quasi-probability distribution at	required
pts_y		y points to evaluate the quasi-probability distribution at; defaults to pts_x	None
g		float, default 2. Scaling factor for a = 0.5 * g * (x + iy). The value of g is related to the value of :math:\hbar in the commutation relation :math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	2
axs		matplotlib axes to plot on (created if None)	None
contour		use contourf if True, otherwise pcolormesh	True
qp_type		type of quasi-probability distribution ("wigner" or "husimi")	WIGNER
cbar_label		label for the cbar (overridden internally based on qp_type)	''
axis_scale_factor		multiplicative scale applied to the axis tick positions and labels	1
plot_cbar		whether to draw a colorbar on each subplot	True
x_ticks		tick positions for the x-axis (auto if None)	None
y_ticks		tick positions for the y-axis (auto if None)	None
z_ticks		tick positions for the colorbar (auto if None)	None
subtitles		subtitles for the subplots; shape must match state.bdims (or the per-frame batch dims when gif=True)	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled subplot grid. Returns a matplotlib.animation.FuncAnimation that auto-renders inline in Jupyter (its _repr_html_ is patched to to_jshtml).	False
gif_params		dict of options for the gif path (ignored if gif=False). Recognized keys: save_path (default None) — if set, save the animation to this path via matplotlib.animation.PillowWriter. interval_ms (default 200) — milliseconds per frame; also derives fps = round(1000 / interval_ms) for the writer. ts (default None) — optional 1D array of timestamps matching the animation-axis length; when set, each frame's suptitle gets a t = … label. batch_animation_axis (default 0) — index into state.bdims selecting which axis becomes the animation axis. The remaining batch dims form the per-frame subplot grid.	None

Returns:

Type	Description
	(axs, im) in the static case, or a FuncAnimation when
	gif=True.

Source code in jaxquantum/core/visualization.py

def plot_qp(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    qp_type=WIGNER,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
    gif=False,
    gif_params=None,
):
    """Plot a quasi-probability distribution (Wigner or Husimi-Q).

    The state may carry an arbitrary number of batch dimensions; they are
    flattened to a 2D ``(rows, cols)`` grid of subplots. With ``gif=True``,
    one batch axis is animated instead and the remaining batch dims form
    the per-frame subplot grid.

    Args:
        state: state with arbitrary number of batch dimensions; result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate the quasi-probability distribution at
        pts_y: y points to evaluate the quasi-probability distribution at;
            defaults to ``pts_x``
        g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``.
            The value of ``g`` is related to the value of :math:`\\hbar` in
            the commutation relation :math:`[x,\,y] = i\\hbar` via
            :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on (created if None)
        contour: use ``contourf`` if True, otherwise ``pcolormesh``
        qp_type: type of quasi-probability distribution
            (``"wigner"`` or ``"husimi"``)
        cbar_label: label for the cbar (overridden internally based on
            ``qp_type``)
        axis_scale_factor: multiplicative scale applied to the axis tick
            positions and labels
        plot_cbar: whether to draw a colorbar on each subplot
        x_ticks: tick positions for the x-axis (auto if None)
        y_ticks: tick positions for the y-axis (auto if None)
        z_ticks: tick positions for the colorbar (auto if None)
        subtitles: subtitles for the subplots; shape must match
            ``state.bdims`` (or the per-frame batch dims when ``gif=True``)
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled subplot grid. Returns a
            ``matplotlib.animation.FuncAnimation`` that auto-renders inline
            in Jupyter (its ``_repr_html_`` is patched to ``to_jshtml``).
        gif_params: dict of options for the gif path (ignored if
            ``gif=False``). Recognized keys:

            - ``save_path`` (default ``None``) — if set, save the animation
              to this path via ``matplotlib.animation.PillowWriter``.
            - ``interval_ms`` (default ``200``) — milliseconds per frame;
              also derives ``fps = round(1000 / interval_ms)`` for the writer.
            - ``ts`` (default ``None``) — optional 1D array of timestamps
              matching the animation-axis length; when set, each frame's
              suptitle gets a ``t = …`` label.
            - ``batch_animation_axis`` (default ``0``) — index into
              ``state.bdims`` selecting which axis becomes the animation
              axis. The remaining batch dims form the per-frame subplot grid.

    Returns:
        ``(axs, im)`` in the static case, or a ``FuncAnimation`` when
        ``gif=True``.
    """
    if pts_y is None:
        pts_y = pts_x
    pts_x = jnp.array(pts_x)
    pts_y = jnp.array(pts_y)

    if len(state.bdims)==1 and state.bdims[0]==1:
        state = state[0]

    if gif:
        return _plot_qp_gif(
            state=state,
            pts_x=pts_x,
            pts_y=pts_y,
            g=g,
            axs=axs,
            contour=contour,
            qp_type=qp_type,
            axis_scale_factor=axis_scale_factor,
            plot_cbar=plot_cbar,
            x_ticks=x_ticks,
            y_ticks=y_ticks,
            z_ticks=z_ticks,
            subtitles=subtitles,
            figtitle=figtitle,
            gif_params=gif_params or {},
        )

    bdims = state.bdims
    added_baxes = 0

    if subtitles is not None and subtitles.shape != bdims:
        raise ValueError(
            f"labels must have same shape as bdims, "
            f"got shapes {subtitles.shape} and {bdims}"
        )

    if len(bdims) == 0:
        bdims = (1,)
        added_baxes += 1
    if len(bdims) == 1:
        bdims = (1, bdims[0])
        added_baxes += 1

    extra_dims = bdims[2:]
    if extra_dims != ():
        state = state.reshape_bdims(
            bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
        )
        if subtitles is not None:
            subtitles = subtitles.reshape(
                bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
            )
        bdims = state.bdims

    if axs is None:
        _, axs = plt.subplots(
            bdims[0],
            bdims[1],
            figsize=(3.3 * bdims[1], 3 * bdims[0]),
            dpi=200,
        )

    if qp_type == WIGNER:
        vmin = -1
        vmax = 1
        scale = np.pi / 2
        cmap = "seismic"
        cbar_label = r"$\mathcal{W}(\alpha)$"
        QP = scale * wigner(state, pts_x, pts_y, g=g)

    elif qp_type == HUSIMI:
        vmin = 0
        vmax = 1
        scale = np.pi
        cmap = "jet"
        cbar_label = r"$\mathcal{Q}(\alpha)$"
        QP = scale * qfunc(state, pts_x, pts_y, g=g)



    for _ in range(added_baxes):
        QP = jnp.array([QP])
        axs = np.array([axs])
        if subtitles is not None:
            subtitles = np.array([subtitles])




    pts_x = pts_x * axis_scale_factor
    pts_y = pts_y * axis_scale_factor

    x_ticks = (
        jnp.linspace(jnp.min(pts_x), jnp.max(pts_x), 5) if x_ticks is None else x_ticks
    )
    y_ticks = (
        jnp.linspace(jnp.min(pts_y), jnp.max(pts_y), 5) if y_ticks is None else y_ticks
    )
    z_ticks = jnp.linspace(vmin, vmax, 3) if z_ticks is None else z_ticks

    im = _render_qp_grid(
        axs,
        QP,
        pts_x,
        pts_y,
        contour=contour,
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        cbar_label=cbar_label,
        plot_cbar=plot_cbar,
        subtitles=subtitles,
        decorate=True,
    )

    fig = axs[bdims[0] - 1, bdims[1] - 1].get_figure()
    fig.tight_layout(w_pad=0.3, h_pad=0.3)
    if figtitle is not None:
        fig.suptitle(figtitle, y=1.04)
    return axs, im

plot_wigner(state, pts_x, pts_y=None, g=2, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot the wigner function of the state.

Thin wrapper around :func:plot_qp with qp_type='wigner'.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
g		float, default 2. Scaling factor for a = 0.5 * g * (x + iy). The value of g is related to the value of :math:\hbar in the commutation relation :math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	2
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled subplot grid. See :func:plot_qp for details.	False
gif_params		dict of options for the gif path. Recognized keys: save_path (default None), interval_ms (default 200), ts (default None — adds a t = … label per frame), batch_animation_axis (default 0).	None

Returns:

Type	Description
	(axs, im) in the static case, or a matplotlib.animation.FuncAnimation
	when gif=True.

Source code in jaxquantum/core/visualization.py

def plot_wigner(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
    gif=False,
    gif_params=None,
):
    """Plot the wigner function of the state.

    Thin wrapper around :func:`plot_qp` with ``qp_type='wigner'``.

    Args:
        state: state with arbitrary number of batch dimensions, result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``.
            The value of ``g`` is related to the value of :math:`\\hbar` in
            the commutation relation :math:`[x,\,y] = i\\hbar` via
            :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled subplot grid. See :func:`plot_qp` for details.
        gif_params: dict of options for the gif path. Recognized keys:
            ``save_path`` (default None), ``interval_ms`` (default 200),
            ``ts`` (default None — adds a ``t = …`` label per frame),
            ``batch_animation_axis`` (default 0).

    Returns:
        ``(axs, im)`` in the static case, or a ``matplotlib.animation.FuncAnimation``
        when ``gif=True``.
    """
    return plot_qp(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        g=g,
        axs=axs,
        contour=contour,
        qp_type=WIGNER,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
        gif=gif,
        gif_params=gif_params,
    )

powm(qarr, n, clip_eigvals=False)

Matrix power of a Qarray.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array.	required
n	Union[int, float]	Exponent. Integer powers use jnp.linalg.matrix_power; float powers diagonalise the matrix.	required
clip_eigvals		When True, clip negative eigenvalues to zero before applying the float power (useful for nearly-PSD matrices).	False

Returns:

Type	Description
Qarray	The n-th matrix power as a Qarray (stays SparseDIA for integer
Qarray	non-negative exponents when the input is SparseDIA).

Raises:

Type	Description
ValueError	If n is a float and the matrix has negative eigenvalues (and clip_eigvals is False).

Source code in jaxquantum/core/qarray.py

def powm(qarr: Qarray, n: Union[int, float], clip_eigvals=False) -> Qarray:
    """Matrix power of a ``Qarray``.

    Args:
        qarr: Input quantum array.
        n: Exponent.  Integer powers use ``jnp.linalg.matrix_power``; float
            powers diagonalise the matrix.
        clip_eigvals: When ``True``, clip negative eigenvalues to zero before
            applying the float power (useful for nearly-PSD matrices).

    Returns:
        The *n*-th matrix power as a ``Qarray`` (stays SparseDIA for integer
        non-negative exponents when the input is SparseDIA).

    Raises:
        ValueError: If *n* is a float and the matrix has negative eigenvalues
            (and *clip_eigvals* is ``False``).
    """
    # SparseDIA fast path: binary exponentiation stays in SparseDIA format.
    if qarr.is_sparse_dia and isinstance(n, int) and n >= 0:
        new_impl = qarr._impl.powm(n)
        return Qarray.create(new_impl.data, dims=qarr.dims, implementation=new_impl.impl_type)

    # Convert to dense for powm
    dense_qarr = qarr.to_dense()

    if isinstance(n, int):
        data_res = jnp.linalg.matrix_power(dense_qarr.data, n)
    else:
        evalues, evectors = jnp.linalg.eig(dense_qarr.data)
        if clip_eigvals:
            evalues = jnp.maximum(evalues, 0)
        else:
            if not (evalues >= 0).all():
                raise ValueError(
                    "Non-integer power of a matrix can only be "
                    "computed if the matrix is positive semi-definite."
                    "Got a matrix with a negative eigenvalue."
                )
        data_res = evectors * jnp.pow(evalues, n) @ jnp.linalg.inv(evectors)

    return Qarray.create(data_res, dims=qarr.dims)

powm_data(data, n)

Integer matrix power of a raw array.

Parameters:

Name	Type	Description	Default
data	Array	Dense square matrix array.	required
n	int	Integer exponent.	required

Returns:

Type	Description
Array	The n-th matrix power.

Source code in jaxquantum/core/qarray.py

def powm_data(data: Array, n: int) -> Array:
    """Integer matrix power of a raw array.

    Args:
        data: Dense square matrix array.
        n: Integer exponent.

    Returns:
        The *n*-th matrix power.
    """
    return jnp.linalg.matrix_power(data, n)

propagator(H, ts, saveat_tlist=None, solver_options=None)

Generate a propagator for a Hamiltonian.

Parameters:

Name	Type	Description	Default
H	Qarray or callable	A Qarray static Hamiltonian OR a function that takes a time argument and returns a Hamiltonian.	required
ts	float or Array	A single time point or an Array of time points.	required
saveat_tlist	Array | None	Times at which to save the propagator.	None
solver_options	SolverOptions | None	Native Diffrax configuration for time-dependent input.	None

Returns:

Type	Description
Qarray	The propagator at each saved time.

Source code in jaxquantum/core/solvers.py

def propagator(
    H: Qarray | Callable[[float], Qarray],
    ts: float | Array,
    saveat_tlist: Array | None = None,
    solver_options: SolverOptions | None = None,
) -> Qarray:
    """Generate a propagator for a Hamiltonian.

    Args:
        H (Qarray or callable):
            A Qarray static Hamiltonian OR
            a function that takes a time argument and returns a Hamiltonian.
        ts (float or Array):
            A single time point or
            an Array of time points.
        saveat_tlist: Times at which to save the propagator.
        solver_options: Native Diffrax configuration for time-dependent input.

    Returns:
        The propagator at each saved time.
    """
    ts_is_scalar = robust_isscalar(ts)
    H_is_qarray = isinstance(H, Qarray)

    if H_is_qarray:
        return (-1j * H * ts).expm()
    else:
        if ts_is_scalar:
            H_first = H(0.0)
            if ts == 0:
                return identity_like(H_first)
            ts = jnp.array([0.0, ts])
        else:
            H_first = H(ts[0])

        basis_states = multi_mode_basis_set(H_first.space_dims)
        results = sesolve(
            H,
            basis_states,
            ts,
            saveat_tlist=saveat_tlist,
            solver_options=solver_options,
        )
        # results.data is (T, M, M): T times, M evolved basis kets (batch), M
        # ket components. Transpose the last two axes so each time slice is a
        # propagator whose columns are the evolved basis states. No squeeze:
        # kets no longer carry a trailing singleton.
        propagators_data = results.data.mT
        return Qarray.create(propagators_data, dims=H_first.space_dims)

ptrace(qarr, indx)

Partial trace over subsystem indx.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array (converted to dense internally).	required
indx		Index of the subsystem to trace out.	required

Returns:

Type	Description
Qarray	Reduced density matrix as a Qarray.

Source code in jaxquantum/core/qarray.py

def ptrace(qarr: Qarray, indx) -> Qarray:
    """Partial trace over subsystem *indx*.

    Args:
        qarr: Input quantum array (converted to dense internally).
        indx: Index of the subsystem to trace out.

    Returns:
        Reduced density matrix as a ``Qarray``.
    """
    # Convert to dense for ptrace
    dense_qarr = qarr.to_dense()
    dense_qarr = ket2dm(dense_qarr)
    rho = dense_qarr.shaped_data
    dims = dense_qarr.dims

    Nq = len(dims[0])

    indxs = [indx, indx + Nq]
    for j in range(Nq):
        if j == indx:
            continue
        indxs.append(j)
        indxs.append(j + Nq)

    bdims = dense_qarr.bdims
    len_bdims = len(bdims)
    bdims_indxs = list(range(len_bdims))
    indxs = bdims_indxs + [j + len_bdims for j in indxs]
    rho = rho.transpose(indxs)

    for j in range(Nq - 1):
        rho = jnp.trace(rho, axis1=2 + len_bdims, axis2=3 + len_bdims)

    return Qarray.create(rho)

qfunc(psi, xvec, yvec, g=2)

Husimi-Q function of a given state vector or density matrix at phase-space points 0.5 * g * (xvec + i*yvec).

Parameters

state : Qarray A state vector or density matrix. This cannot have tensor-product structure.

xvec, yvec : array_like x- and y-coordinates at which to calculate the Husimi-Q function.

float, default: 2

Scaling factor for a = 0.5 * g * (x + iy). The value of g is related to the value of :math:\hbar in the commutation relation :math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.

Returns

jnp.ndarray Values representing the Husimi-Q function calculated over the specified range [xvec, yvec].

Source code in jaxquantum/core/qp_distributions.py

def qfunc(psi, xvec, yvec, g=2):
    r"""
    Husimi-Q function of a given state vector or density matrix at phase-space
    points ``0.5 * g * (xvec + i*yvec)``.

    Parameters
    ----------
    state : Qarray
        A state vector or density matrix. This cannot have tensor-product
        structure.

    xvec, yvec : array_like
        x- and y-coordinates at which to calculate the Husimi-Q function.

    g : float, default: 2
        Scaling factor for ``a = 0.5 * g * (x + iy)``.  The value of `g` is
        related to the value of :math:`\hbar` in the commutation relation
        :math:`[x,\,y] = i\hbar` via :math:`\hbar=2/g^2`.

    Returns
    -------
    jnp.ndarray
        Values representing the Husimi-Q function calculated over the specified
        range ``[xvec, yvec]``.

    """

    alpha_grid, prefactor = _qfunc_coherent_grid(xvec, yvec, g)

    if psi.is_vec():
        psi = psi.to_ket()

        def _compute_qfunc(psi, alpha_grid, prefactor, g):
            out = _qfunc_iterative_single(psi, alpha_grid, prefactor, g)
            out /= jnp.pi
            return out
    else:

        def _compute_qfunc(psi, alpha_grid, prefactor, g):
            values, vectors = jnp.linalg.eigh(psi)
            vectors = vectors.T
            chunk_size = min(8, values.shape[0])
            chunks = (values.shape[0] + chunk_size - 1) // chunk_size
            padding = chunks * chunk_size - values.shape[0]
            values = jnp.pad(values, (0, padding))
            vectors = jnp.pad(vectors, ((0, padding), (0, 0)))

            def add_chunk(index, total):
                start = index * chunk_size
                chunk_values = lax.dynamic_slice_in_dim(
                    values, start, chunk_size
                )
                chunk_vectors = lax.dynamic_slice_in_dim(
                    vectors, start, chunk_size
                )
                components = vmap(
                    lambda vector: _qfunc_iterative_single(
                        vector, alpha_grid, prefactor, g
                    )
                )(chunk_vectors)
                return total + jnp.tensordot(chunk_values, components, axes=1)

            out = lax.fori_loop(
                0,
                chunks,
                add_chunk,
                jnp.zeros_like(alpha_grid.real),
            )
            out /= jnp.pi

            return out

    # Use the Qarray batch dims rather than data.shape[:-2]: kets now have a
    # single trailing space axis, so shape[:-2] would under-count batch dims.
    bdims = psi.bdims
    psi = psi.data

    vmapped_compute_qfunc = [_compute_qfunc]

    for _ in bdims:
        vmapped_compute_qfunc.append(
            vmap(
                vmapped_compute_qfunc[-1],
                in_axes=(0, None, None, None),
                out_axes=0,
            )
        )
    return vmapped_compute_qfunc[-1](psi, alpha_grid, prefactor, g)

qt2jqt(qt_obj, dtype=jnp.complex128)

QuTiP state -> Qarray.

Parameters:

Name	Type	Description	Default
qt_obj		QuTiP state.	required
dtype		JAX dtype.	complex128

Returns:

Type	Description
	Qarray.

Source code in jaxquantum/core/conversions.py

def qt2jqt(qt_obj, dtype=jnp.complex128):
    """QuTiP state -> Qarray.

    Args:
        qt_obj: QuTiP state.
        dtype: JAX dtype.

    Returns:
        Qarray.
    """
    if isinstance(qt_obj, Qarray) or qt_obj is None:
        return qt_obj
    return Qarray.create(jnp.array(qt_obj.full(), dtype=dtype), dims=qt_obj.dims)

qubit_rotation(theta, nx, ny, nz)

Single qubit rotation.

Parameters:

Name	Type	Description	Default
theta	float	rotation angle.	required
nx		rotation axis x component.	required
ny		rotation axis y component.	required
nz		rotation axis z component.	required

Returns:

Type	Description
Qarray	Single qubit rotation operator.

Source code in jaxquantum/core/operators.py

def qubit_rotation(theta: float, nx, ny, nz) -> Qarray:
    """Single qubit rotation.

    Args:
        theta: rotation angle.
        nx: rotation axis x component.
        ny: rotation axis y component.
        nz: rotation axis z component.

    Returns:
        Single qubit rotation operator.
    """
    return jnp.cos(theta / 2) * identity(2) - 1j * jnp.sin(theta / 2) * (
        nx * sigmax() + ny * sigmay() + nz * sigmaz()
    )

robust_asarray(data)

Convert data to a JAX array, leaving sparse BCOO and SparseDiaData untouched.

Parameters:

Name	Type	Description	Default
data		Input data — any array-like, sparse.BCOO, or SparseDiaData.	required

Returns:

Type	Description
Union[Array, BCOO]	A jax.Array, sparse.BCOO, or SparseDiaData.

Source code in jaxquantum/core/qarray.py

def robust_asarray(data) -> Union[Array, sparse.BCOO]:
    """Convert *data* to a JAX array, leaving sparse BCOO and SparseDiaData untouched.

    Args:
        data: Input data — any array-like, ``sparse.BCOO``, or ``SparseDiaData``.

    Returns:
        A ``jax.Array``, ``sparse.BCOO``, or ``SparseDiaData``.
    """
    if isinstance(data, sparse.BCOO):
        return data
    # SparseDiaData has a ``_is_sparse_dia`` marker; pass it through unchanged
    if getattr(data, "_is_sparse_dia", False):
        return data
    return jnp.asarray(data)

sesolve(H, rho0, tlist, saveat_tlist=None, solver_options=None)

Solve a Schrödinger equation and return the saved states.

Parameters:

Name	Type	Description	Default
H	Qarray | Callable[[float], Qarray]	Static Hamiltonian or callable H(t).	required
rho0	Qarray	Initial ket.	required
tlist	Array	Integration interval; also the default save times.	required
saveat_tlist	Array | None	Save times. An empty array saves only the final state.	None
solver_options	SolverOptions | None	Native Diffrax configuration.	None

Returns:

Type	Description
Qarray	Saved kets as a batched Qarray.

See Also

:func:sesolve_result returns the complete Diffrax solution.

Source code in jaxquantum/core/solvers.py

def sesolve(
    H: Qarray | Callable[[float], Qarray],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Array | None = None,
    solver_options: SolverOptions | None = None,
) -> Qarray:
    """Solve a Schrödinger equation and return the saved states.

    Args:
        H: Static Hamiltonian or callable ``H(t)``.
        rho0: Initial ket.
        tlist: Integration interval; also the default save times.
        saveat_tlist: Save times. An empty array saves only the final state.
        solver_options: Native Diffrax configuration.

    Returns:
        Saved kets as a batched ``Qarray``.

    See Also:
        :func:`sesolve_result` returns the complete Diffrax solution.
    """
    solution = sesolve_result(
        H,
        rho0,
        tlist,
        saveat_tlist=saveat_tlist,
        solver_options=solver_options,
    )
    return Qarray._from_impl(DenseImpl._make(solution.ys), rho0.qdims)

sesolve_result(H, rho0, tlist, saveat_tlist=None, solver_options=None)

Solve a Schrödinger equation and return its Diffrax solution.

Use this form for solver statistics, events, dense interpolation, custom SaveAt functions, or continuation state.

Source code in jaxquantum/core/solvers.py

def sesolve_result(
    H: Qarray | Callable[[float], Qarray],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Array | None = None,
    solver_options: SolverOptions | None = None,
) -> diffrax.Solution:
    """Solve a Schrödinger equation and return its Diffrax solution.

    Use this form for solver statistics, events, dense interpolation, custom
    ``SaveAt`` functions, or continuation state.
    """
    if rho0.qtype == Qtypes.oper:
        raise ValueError("Use mesolve() for an initial density matrix.")

    state = rho0.to_ket().to_dense()

    if robust_isscalar(H):
        H = H * identity_like(state)

    if isinstance(H, Qarray):
        H_data = lambda t: H.data
    else:
        H_data = lambda t: H(t).data

    return _sesolve_result_data(
        H_data,
        state.data,
        tlist,
        saveat_tlist,
        solver_options=solver_options,
    )

sigmam(implementation=QarrayImplType.DENSE)

σ-

Returns:

Type	Description
Qarray	σ- Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmam(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """σ-

    Returns:
        σ- Pauli Operator
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        diags = jnp.array([[1.0, 0.0]])
        return _make_sparsedia(offsets=(-1,), diags=diags)
    return Qarray.create(jnp.array([[0.0, 0.0], [1.0, 0.0]]), implementation=implementation)

sigmap(implementation=QarrayImplType.DENSE)

σ+

Returns:

Type	Description
Qarray	σ+ Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmap(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """σ+

    Returns:
        σ+ Pauli Operator
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        diags = jnp.array([[0.0, 1.0]])
        return _make_sparsedia(offsets=(1,), diags=diags)
    return Qarray.create(jnp.array([[0.0, 1.0], [0.0, 0.0]]), implementation=implementation)

sigmax(implementation=QarrayImplType.DENSE)

σx

Parameters:

Name	Type	Description	Default
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	σx Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmax(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """σx

    Args:
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        σx Pauli Operator
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # Offset -1: valid at [0:1] → diag[0] = A[1,0] = 1.0, diag[1] = 0 (trailing zero)
        # Offset +1: valid at [1:]  → diag[0] = 0 (leading zero), diag[1] = A[0,1] = 1.0
        diags = jnp.array([[1.0, 0.0], [0.0, 1.0]])
        return _make_sparsedia(offsets=(-1, 1), diags=diags)
    return Qarray.create(jnp.array([[0.0, 1.0], [1.0, 0.0]]), implementation=implementation)

sigmay(implementation=QarrayImplType.DENSE)

σy

Returns:

Type	Description
Qarray	σy Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmay(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """σy

    Returns:
        σy Pauli Operator
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        diags = jnp.array([[1.0j, 0.0], [0.0, -1.0j]])
        return _make_sparsedia(offsets=(-1, 1), diags=diags)
    return Qarray.create(jnp.array([[0.0, -1.0j], [1.0j, 0.0]]), implementation=implementation)

sigmaz(implementation=QarrayImplType.DENSE)

σz

Returns:

Type	Description
Qarray	σz Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmaz(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """σz

    Returns:
        σz Pauli Operator
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        diags = jnp.array([[1.0, -1.0]])
        return _make_sparsedia(offsets=(0,), diags=diags)
    return Qarray.create(jnp.array([[1.0, 0.0], [0.0, -1.0]]), implementation=implementation)

sinm(qarr)

Matrix sine of a Qarray.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array (converted to dense internally).	required

Returns:

Type	Description
Qarray	A dense Qarray containing the matrix sine.

Source code in jaxquantum/core/qarray.py

def sinm(qarr: Qarray) -> Qarray:
    """Matrix sine of a ``Qarray``.

    Args:
        qarr: Input quantum array (converted to dense internally).

    Returns:
        A dense ``Qarray`` containing the matrix sine.
    """
    dims = qarr.dims
    # Convert to dense for sinm
    dense_data = qarr.to_dense().data
    data = sinm_data(dense_data)
    return Qarray.create(data, dims=dims)

sinm_data(data, **kwargs)

Matrix sine of a raw array.

Parameters:

Name	Type	Description	Default
data	Array	Dense matrix array.	required
**kwargs		Unused; kept for API consistency.	{}

Returns:

Type	Description
Array	The matrix sine computed as (expm(i*A) - expm(-i*A)) / (2i).

Source code in jaxquantum/core/qarray.py

def sinm_data(data: Array, **kwargs) -> Array:
    """Matrix sine of a raw array.

    Args:
        data: Dense matrix array.
        **kwargs: Unused; kept for API consistency.

    Returns:
        The matrix sine computed as ``(expm(i*A) - expm(-i*A)) / (2i)``.
    """
    return (expm_data(1j * data) - expm_data(-1j * data)) / (2j)

solve(f, y0, tlist, saveat_tlist=None, args=None, solver_options=None)

Solve an ODE using native Diffrax configuration from SolverOptions.

Source code in jaxquantum/core/solvers.py

def solve(
    f: Callable,
    y0: Array,
    tlist: Array,
    saveat_tlist: Array | None = None,
    args: Any = None,
    solver_options: SolverOptions | None = None,
) -> diffrax.Solution:
    """Solve an ODE using native Diffrax configuration from ``SolverOptions``."""
    options = SolverOptions() if solver_options is None else solver_options
    if _uses_legacy_options(options):
        warnings.warn(
            "String and boolean SolverOptions values are deprecated; use native "
            "objects such as solver=diffrax.Tsit5(), "
            "stepsize_controller=diffrax.PIDController(...), and "
            "progress_meter='default' or None.",
            FutureWarning,
            stacklevel=2,
        )
    kwargs = {
        "saveat": _resolve_saveat(options, tlist, saveat_tlist),
        "stepsize_controller": _resolve_stepsize_controller(options),
        "args": args,
        "max_steps": options.max_steps,
        "throw": options.throw,
    }
    optional = {
        "adjoint": options.adjoint,
        "event": options.event,
        "progress_meter": _resolve_progress_meter(options.progress_meter),
        "solver_state": options.solver_state,
        "controller_state": options.controller_state,
        "made_jump": options.made_jump,
    }
    kwargs.update(
        (name, value) for name, value in optional.items() if value is not None
    )

    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="Complex dtype support in Diffrax",
            category=UserWarning,
        )
        return diffrax.diffeqsolve(
            diffrax.ODETerm(f),
            _resolve_solver(options.solver),
            t0=tlist[0],
            t1=tlist[-1],
            dt0=_resolve_dt0(options, tlist),
            y0=y0,
            **kwargs,
        )

squeeze(N, z)

Single-mode Squeezing operator.

Parameters:

Name	Type	Description	Default
N		Hilbert Space Size	required
z		squeezing parameter	required

Returns:

Type	Description
	Sqeezing operator

Source code in jaxquantum/core/operators.py

def squeeze(N, z):
    """Single-mode Squeezing operator.


    Args:
        N: Hilbert Space Size
        z: squeezing parameter

    Returns:
        Sqeezing operator
    """

    a = destroy(N)
    op = (1 / 2.0) * jnp.conj(z) * (a @ a) - (1 / 2.0) * z * (a.dag() @ a.dag())
    return op.expm()

tensor(*args, **kwargs)

Tensor (Kronecker) product of two or more Qarray objects.

Parameters:

Name	Type	Description	Default
*args		Qarray objects to tensor together (left to right).	()
**kwargs		Optional keyword arguments. Pass parallel=True to use an einsum-based batched outer product instead of jnp.kron.	{}

Returns:

Type	Description
Qarray	The tensor product as a Qarray. The output implementation is
Qarray	determined by the highest PROMOTION_ORDER among the inputs: all-sparse
Qarray	inputs → sparse output; any dense input → dense output. This holds for
Qarray	both parallel=True and parallel=False.

Note

parallel=True uses an einsum-based batched outer product. The einsum is always computed on dense data for efficiency, but the result is then wrapped in the appropriate backend (sparse when all inputs are sparse, dense otherwise). For the default (parallel=False) path each backend's kron method is used directly.

Source code in jaxquantum/core/qarray.py

def tensor(*args, **kwargs) -> Qarray:
    """Tensor (Kronecker) product of two or more ``Qarray`` objects.

    Args:
        *args: ``Qarray`` objects to tensor together (left to right).
        **kwargs: Optional keyword arguments.  Pass ``parallel=True`` to use
            an einsum-based batched outer product instead of ``jnp.kron``.

    Returns:
        The tensor product as a ``Qarray``.  The output implementation is
        determined by the highest ``PROMOTION_ORDER`` among the inputs: all-sparse
        inputs → sparse output; any dense input → dense output.  This holds for
        both ``parallel=True`` and ``parallel=False``.

    Note:
        ``parallel=True`` uses an einsum-based batched outer product.  The
        einsum is always computed on dense data for efficiency, but the result
        is then wrapped in the appropriate backend (sparse when all inputs are
        sparse, dense otherwise).  For the default (``parallel=False``) path
        each backend's ``kron`` method is used directly.
    """
    parallel = kwargs.pop("parallel", False)

    if parallel:
        # Determine target implementation: highest PROMOTION_ORDER wins.
        # All-sparse → sparse; any dense input → dense (same rule as non-parallel).
        target_impl_type = max(
            (arg.impl_type for arg in args),
            key=lambda t: t.get_impl_class().PROMOTION_ORDER,
        )
        # Einsum-based batched outer product (computed on dense data).
        dense_args = [arg.to_dense() for arg in args]
        # Vectors (kets/bras) carry one trailing space axis; operators carry two.
        # Tensoring requires a consistent qtype across args (mixed is invalid).
        is_vec = dense_args[0].qtype in (Qtypes.ket, Qtypes.bra)
        n_space = 1 if is_vec else 2
        data = dense_args[0].data
        dims_0 = dense_args[0].dims[0]
        dims_1 = dense_args[0].dims[1]
        for arg in dense_args[1:]:
            a, b = data, arg.data
            ba, bb = a.shape[:-n_space], b.shape[:-n_space]
            if len(ba) > len(bb):
                batch_dim = ba
            elif len(ba) == len(bb):
                batch_dim = ba if prod(ba) >= prod(bb) else bb
            else:
                batch_dim = bb

            if is_vec:
                # (..., N) ⊗ (..., M) -> (..., N*M)
                data = jnp.einsum("...i,...j->...ij", a, b).reshape(
                    *batch_dim, a.shape[-1] * b.shape[-1]
                )
            else:
                # (..., M, N) ⊗ (..., K, L) -> (..., M*K, N*L)
                data = jnp.einsum("...ij,...kl->...ikjl", a, b).reshape(
                    *batch_dim, a.shape[-2] * b.shape[-2], -1
                )
            dims_0 = dims_0 + arg.dims[0]
            dims_1 = dims_1 + arg.dims[1]
        impl = target_impl_type.get_impl_class().from_data(data)
        return Qarray._from_impl(impl, Qdims((dims_0, dims_1)))

    # Non-parallel: delegate to each impl's kron method.
    # All-sparse inputs stay sparse; mixed inputs promote to dense via _coerce.
    current_impl = args[0]._impl
    dims_0 = args[0].dims[0]
    dims_1 = args[0].dims[1]
    for arg in args[1:]:
        current_impl = current_impl.kron(arg._impl)
        dims_0 = dims_0 + arg.dims[0]
        dims_1 = dims_1 + arg.dims[1]
    return Qarray._from_impl(current_impl, Qdims((dims_0, dims_1)))

tensor_basis(single_basis, n)

Construct n-fold tensor product basis from a single-system basis.

Parameters:

Name	Type	Description	Default
single_basis	Qarray	The single-system operator basis as a Qarray.	required
n	int	Number of tensor copies to construct.	required

Returns:

Type	Description
Qarray	Qarray containing the n-fold tensor product basis operators.
Qarray	The resulting basis has b^n elements where b is the number
Qarray	of operators in the single-system basis.

Source code in jaxquantum/core/measurements.py

def tensor_basis(single_basis: Qarray, n: int) -> Qarray:
    """Construct n-fold tensor product basis from a single-system basis.

    Args:
        single_basis: The single-system operator basis as a Qarray.
        n: Number of tensor copies to construct.

    Returns:
        Qarray containing the n-fold tensor product basis operators.
        The resulting basis has b^n elements where b is the number
        of operators in the single-system basis.
    """

    dims = single_basis.dims

    single_basis = single_basis.data
    b, _, _ = single_basis.shape
    indices = jnp.stack(jnp.meshgrid(*[jnp.arange(b)] * n, indexing="ij"),
                        axis=-1).reshape(-1, n)  # shape (b^n, n)

    # Select the operators based on indices: shape (b^n, n, d, d)
    selected = single_basis[indices]  # shape: (b^n, n, d, d)

    # Vectorized Kronecker products
    full_basis = vmap(lambda ops: reduce(jnp.kron, ops))(selected)

    new_dims = tuple(tuple(x**n for x in row) for row in dims)

    return Qarray.create(full_basis, dims=new_dims, bdims=(b**n,))

thermal_dm(N, n)

Thermal state.

Parameters:

Name	Type	Description	Default
N	int	Hilbert Space Size.	required
n	float	average photon number.	required

Return

Thermal state.

Source code in jaxquantum/core/operators.py

def thermal_dm(N: int, n: float) -> Qarray:
    """Thermal state.

    Args:
        N: Hilbert Space Size.
        n: average photon number.

    Return:
        Thermal state.
    """

    beta = jnp.log(1 + 1 / n)

    return Qarray.create(
        jnp.where(
            jnp.isposinf(beta),
            basis(N, 0).to_dm().data,
            jnp.diag(jnp.exp(-beta * jnp.linspace(0, N - 1, N))),
        )
    ).unit()

to_ket(qarr)

Convert qarr to a ket.

Parameters:

Name	Type	Description	Default
qarr	Qarray	A ket (returned as-is) or bra (conjugate-transposed).	required

Returns:

Type	Description
Qarray	The ket form of qarr.

Raises:

Type	Description
ValueError	If qarr is an operator.

Source code in jaxquantum/core/qarray.py

def to_ket(qarr: Qarray) -> Qarray:
    """Convert *qarr* to a ket.

    Args:
        qarr: A ket (returned as-is) or bra (conjugate-transposed).

    Returns:
        The ket form of *qarr*.

    Raises:
        ValueError: If *qarr* is an operator.
    """
    if qarr.qtype == Qtypes.ket:
        return qarr
    elif qarr.qtype == Qtypes.bra:
        return qarr.dag()
    else:
        raise ValueError("Can only get ket from a ket or bra.")

tr(qarr, **kwargs)

Full trace of qarr.

For sparse Qarray objects the trace is computed natively on the BCOO data using a masked scatter — no densification. Custom axis arguments are ignored for sparse (the last two dimensions are always the matrix dimensions in jaxquantum's convention).

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array.	required
**kwargs		Forwarded to jnp.trace for dense arrays (e.g. axis1, axis2).	{}

Returns:

Type	Description
Array	The trace as a scalar (or batched array of scalars).

Source code in jaxquantum/core/qarray.py

def tr(qarr: Qarray, **kwargs) -> Array:
    """Full trace of *qarr*.

    For sparse ``Qarray`` objects the trace is computed natively on the BCOO
    data using a masked scatter — no densification.  Custom axis arguments
    are ignored for sparse (the last two dimensions are always the matrix
    dimensions in jaxquantum's convention).

    Args:
        qarr: Input quantum array.
        **kwargs: Forwarded to ``jnp.trace`` for dense arrays (e.g.
            ``axis1``, ``axis2``).

    Returns:
        The trace as a scalar (or batched array of scalars).
    """
    if qarr.is_sparse_bcoo:
        return qarr._impl.trace()
    if qarr.is_sparse_dia:
        return qarr._impl.trace()
    axis1 = kwargs.get("axis1", -2)
    axis2 = kwargs.get("axis2", -1)
    return jnp.trace(qarr.data, axis1=axis1, axis2=axis2, **kwargs)

trace(qarr, **kwargs)

Full trace (alias for :func:tr).

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array.	required
**kwargs		Forwarded to :func:tr.	{}

Returns:

Type	Description
Array	The trace as a scalar (or batched array of scalars).

Source code in jaxquantum/core/qarray.py

def trace(qarr: Qarray, **kwargs) -> Array:
    """Full trace (alias for :func:`tr`).

    Args:
        qarr: Input quantum array.
        **kwargs: Forwarded to :func:`tr`.

    Returns:
        The trace as a scalar (or batched array of scalars).
    """
    return tr(qarr, **kwargs)

transpose(qarr, indices)

Transpose subsystem indices of the quantum array.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array.	required
indices	List[int]	New ordering of subsystem indices.	required

Returns:

Type	Description
Qarray	Transposed Qarray (converted to dense first).

Source code in jaxquantum/core/qarray.py

def transpose(qarr: Qarray, indices: List[int]) -> Qarray:
    """Transpose subsystem indices of the quantum array.

    Args:
        qarr: Input quantum array.
        indices: New ordering of subsystem indices.

    Returns:
        Transposed ``Qarray`` (converted to dense first).
    """

    qarr = qarr.to_dense()

    indices = list(indices)

    shaped_data = qarr.shaped_data
    dims = qarr.dims
    bdims_indxs = list(range(len(qarr.bdims)))

    reshape_indices = indices + [j + len(dims[0]) for j in indices]
    reshape_indices = bdims_indxs + [j + len(bdims_indxs) for j in reshape_indices]

    shaped_data = shaped_data.transpose(reshape_indices)
    new_dims = (
        tuple([dims[0][j] for j in indices]),
        tuple([dims[1][j] for j in indices]),
    )

    full_dims = prod(dims[0])
    full_data = shaped_data.reshape(*qarr.bdims, full_dims, -1)

    # Preserve implementation type
    implementation = qarr.impl_type
    return Qarray.create(full_data, dims=new_dims, implementation=implementation)

unit(qarr)

Normalize qarr to unit norm.

Parameters:

Name	Type	Description	Default
qarr	Qarray	Input quantum array.	required

Returns:

Type	Description
Qarray	Normalized quantum array.

Source code in jaxquantum/core/qarray.py

def unit(qarr: Qarray) -> Qarray:
    """Normalize *qarr* to unit norm.

    Args:
        qarr: Input quantum array.

    Returns:
        Normalized quantum array.
    """
    return qarr / qarr.norm()

wigner(psi, xvec, yvec, method='clenshaw', g=2)

Wigner function for a state vector or density matrix at points xvec + i * yvec.

Parameters

Qarray

A state vector or density matrix.

array_like

x-coordinates at which to calculate the Wigner function.

array_like

y-coordinates at which to calculate the Wigner function.

float, default: 2

Scaling factor for a = 0.5 * g * (x + iy), default g = 2. The value of g is related to the value of hbar in the commutation relation [x, y] = i * hbar via hbar=2/g^2.

string {'clenshaw', 'iterative', 'laguerre', 'fft'}, default: 'clenshaw'

Only 'clenshaw' is currently supported. Select method 'clenshaw' 'iterative', 'laguerre', or 'fft', where 'clenshaw' and 'iterative' use an iterative method to evaluate the Wigner functions for density matrices :math:|m><n|, while 'laguerre' uses the Laguerre polynomials in scipy for the same task. The 'fft' method evaluates the Fourier transform of the density matrix. The 'iterative' method is default, and in general recommended, but the 'laguerre' method is more efficient for very sparse density matrices (e.g., superpositions of Fock states in a large Hilbert space). The 'clenshaw' method is the preferred method for dealing with density matrices that have a large number of excitations (>~50). 'clenshaw' is a fast and numerically stable method.

Returns

array

Values representing the Wigner function calculated over the specified range [xvec,yvec].

References

Ulf Leonhardt, Measuring the Quantum State of Light, (Cambridge University Press, 1997)

Source code in jaxquantum/core/qp_distributions.py

def wigner(psi, xvec, yvec, method="clenshaw", g=2):
    """Wigner function for a state vector or density matrix at points
    `xvec + i * yvec`.

    Parameters
    ----------

    state : Qarray
        A state vector or density matrix.

    xvec : array_like
        x-coordinates at which to calculate the Wigner function.

    yvec : array_like
        y-coordinates at which to calculate the Wigner function.

    g : float, default: 2
        Scaling factor for `a = 0.5 * g * (x + iy)`, default `g = 2`.
        The value of `g` is related to the value of `hbar` in the commutation
        relation `[x, y] = i * hbar` via `hbar=2/g^2`.

    method : string {'clenshaw', 'iterative', 'laguerre', 'fft'}, default: 'clenshaw'
        Only 'clenshaw' is currently supported.
        Select method 'clenshaw' 'iterative', 'laguerre', or 'fft', where 'clenshaw'
        and 'iterative' use an iterative method to evaluate the Wigner functions for density
        matrices :math:`|m><n|`, while 'laguerre' uses the Laguerre polynomials
        in scipy for the same task. The 'fft' method evaluates the Fourier
        transform of the density matrix. The 'iterative' method is default, and
        in general recommended, but the 'laguerre' method is more efficient for
        very sparse density matrices (e.g., superpositions of Fock states in a
        large Hilbert space). The 'clenshaw' method is the preferred method for
        dealing with density matrices that have a large number of excitations
        (>~50). 'clenshaw' is a fast and numerically stable method.

    Returns
    -------

    W : array
        Values representing the Wigner function calculated over the specified
        range [xvec,yvec].


    References
    ----------

    Ulf Leonhardt,
    Measuring the Quantum State of Light, (Cambridge University Press, 1997)

    """

    if not (psi.is_vec() or psi.is_dm()):
        raise TypeError("Input state is not a valid operator.")

    if method == "fft":
        raise NotImplementedError("Only the 'clenshaw' method is implemented.")

    if method == "iterative":
        raise NotImplementedError("Only the 'clenshaw' method is implemented.")

    elif method == "laguerre":
        raise NotImplementedError("Only the 'clenshaw' method is implemented.")

    elif method == "clenshaw":
        rho = psi.to_dm()
        bdims = rho.bdims
        rho = rho.data

        vmapped_wigner_clenshaw = [_wigner_clenshaw]

        for _ in bdims:
            vmapped_wigner_clenshaw.append(
                vmap(
                    vmapped_wigner_clenshaw[-1],
                    in_axes=(0, None, None, None),
                    out_axes=0,
                )
            )
        return vmapped_wigner_clenshaw[-1](rho, xvec, yvec, g)

    else:
        raise TypeError("method must be either 'iterative', 'laguerre', or 'fft'.")