base

Base.

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))

Device

Bases: ABC

Source code in jaxquantum/devices/base/base.py

@struct.dataclass
class Device(ABC):
    DEFAULT_BASIS = BasisTypes.fock
    DEFAULT_HAMILTONIAN = HamiltonianTypes.full

    N: int = struct.field(pytree_node=False)
    N_pre_diag: int = struct.field(pytree_node=False)
    params: Dict[str, Any]
    _label: int = struct.field(pytree_node=False)
    _basis: BasisTypes = struct.field(pytree_node=False)
    _hamiltonian: HamiltonianTypes = struct.field(pytree_node=False)

    @classmethod
    def param_validation(cls, N, N_pre_diag, params, hamiltonian, basis):
        """This can be overridden by subclasses."""
        pass

    @classmethod
    def create(
        cls,
        N,
        params,
        label=0,
        N_pre_diag=None,
        use_linear=False,
        hamiltonian: HamiltonianTypes = None,
        basis: BasisTypes = None,
    ):
        """Create a device.

        Args:
            N (int): dimension of Hilbert space.
            params (dict): parameters of the device.
            label (int, optional): label for the device. Defaults to 0. This is useful when you have multiple of the same device type in the same system.
            N_pre_diag (int, optional): dimension of Hilbert space before diagonalization. Defaults to None, in which case it is set to N. This must be greater than or rqual to N.
            use_linear (bool): whether to use the linearized device. Defaults to False. This will override the hamiltonian keyword argument. This is a bit redundant with hamiltonian, but it is kept for backwards compatibility.
            hamiltonian (HamiltonianTypes, optional): type of Hamiltonian. Defaults to None, in which case the full hamiltonian is used.
            basis (BasisTypes, optional): type of basis. Defaults to None, in which case the fock basis is used.
        """

        if N_pre_diag is None:
            N_pre_diag = N

        assert N_pre_diag >= N, "N_pre_diag must be greater than or equal to N."

        _basis = basis if basis is not None else cls.DEFAULT_BASIS
        _hamiltonian = (
            hamiltonian if hamiltonian is not None else cls.DEFAULT_HAMILTONIAN
        )

        if use_linear:
            _hamiltonian = HamiltonianTypes.linear

        cls.param_validation(N, N_pre_diag, params, _hamiltonian, _basis)

        return cls(N, N_pre_diag, params, label, _basis, _hamiltonian)

    @property
    def basis(self):
        return self._basis

    @property
    def hamiltonian(self):
        return self._hamiltonian

    @property
    def label(self):
        return self.__class__.__name__ + str(self._label)

    @property
    def linear_ops(self):
        return self.common_ops()

    @property
    def original_ops(self):
        return self.common_ops()

    @property
    def ops(self):
        return self.full_ops()

    @abstractmethod
    def common_ops(self) -> Dict[str, Qarray]:
        """Set up common ops in the specified basis."""

    @abstractmethod
    def get_linear_frequency(self):
        """Get frequency of linear terms."""

    @abstractmethod
    def get_H_linear(self):
        """Return linear terms in H."""

    @abstractmethod
    def get_H_full(self):
        """Return full H."""

    def get_H(self):
        """
        Return the Hamiltonian truncated in its eigenbasis.
        """
        eig_systems = self.eig_systems
        values = eig_systems["vals"][..., : self.N]
        dtype = eig_systems["vecs"].dtype
        data = values.astype(dtype)[..., :, None] * jnp.eye(self.N, dtype=dtype)
        qdims = Qdims(((self.N,), (self.N,)))
        return Qarray._from_impl(DenseImpl._make(data), qdims)

    def _get_H_in_original_basis(self):
        """This returns the Hamiltonian in the original specified basis. This can be overridden by subclasses."""

        if self.hamiltonian == HamiltonianTypes.linear:
            return self.get_H_linear()
        elif self.hamiltonian == HamiltonianTypes.full:
            return self.get_H_full()

    def _calculate_eig_systems(self):
        return jnp.linalg.eigh(self._get_H_in_original_basis().data)

    @property
    def eig_systems(self):
        values, vectors = self._calculate_eig_systems()
        return {"vals": values, "vecs": vectors}

    def get_op_in_H_eigenbasis(self, op: Qarray):
        evecs = self.eig_systems["vecs"][..., :, : self.N]
        dims = [[self.N], [self.N]]
        return get_op_in_new_basis(op, evecs, dims)

    def get_op_data_in_H_eigenbasis(self, op: Array):
        evecs = self.eig_systems["vecs"][..., :, : self.N]
        return get_op_data_in_new_basis(op, evecs)

    def get_vec_in_H_eigenbasis(self, vec: Qarray):
        evecs = self.eig_systems["vecs"][..., :, : self.N]
        if vec.qtype == Qtypes.ket:
            dims = [[self.N], [1]]
        else:
            dims = [[1], [self.N]]
        return get_vec_in_new_basis(vec, evecs, dims)

    def get_vec_data_in_H_eigenbasis(self, vec: Array):
        evecs = self.eig_systems["vecs"][..., :, : self.N]
        return get_vec_data_in_new_basis(vec, evecs)

    def full_ops(self):
        linear_ops = self.linear_ops
        vectors = self.eig_systems["vecs"][..., :, : self.N]
        dims = ((self.N,), (self.N,))
        return {
            name: get_op_in_new_basis(op, vectors, dims)
            for name, op in linear_ops.items()
        }

common_ops() abstractmethod

Set up common ops in the specified basis.

Source code in jaxquantum/devices/base/base.py

@abstractmethod
def common_ops(self) -> Dict[str, Qarray]:
    """Set up common ops in the specified basis."""

create(N, params, label=0, N_pre_diag=None, use_linear=False, hamiltonian=None, basis=None) classmethod

Create a device.

Parameters:

Name	Type	Description	Default
N	int	dimension of Hilbert space.	required
params	dict	parameters of the device.	required
label	int	label for the device. Defaults to 0. This is useful when you have multiple of the same device type in the same system.	0
N_pre_diag	int	dimension of Hilbert space before diagonalization. Defaults to None, in which case it is set to N. This must be greater than or rqual to N.	None
use_linear	bool	whether to use the linearized device. Defaults to False. This will override the hamiltonian keyword argument. This is a bit redundant with hamiltonian, but it is kept for backwards compatibility.	False
hamiltonian	HamiltonianTypes	type of Hamiltonian. Defaults to None, in which case the full hamiltonian is used.	None
basis	BasisTypes	type of basis. Defaults to None, in which case the fock basis is used.	None

Source code in jaxquantum/devices/base/base.py

@classmethod
def create(
    cls,
    N,
    params,
    label=0,
    N_pre_diag=None,
    use_linear=False,
    hamiltonian: HamiltonianTypes = None,
    basis: BasisTypes = None,
):
    """Create a device.

    Args:
        N (int): dimension of Hilbert space.
        params (dict): parameters of the device.
        label (int, optional): label for the device. Defaults to 0. This is useful when you have multiple of the same device type in the same system.
        N_pre_diag (int, optional): dimension of Hilbert space before diagonalization. Defaults to None, in which case it is set to N. This must be greater than or rqual to N.
        use_linear (bool): whether to use the linearized device. Defaults to False. This will override the hamiltonian keyword argument. This is a bit redundant with hamiltonian, but it is kept for backwards compatibility.
        hamiltonian (HamiltonianTypes, optional): type of Hamiltonian. Defaults to None, in which case the full hamiltonian is used.
        basis (BasisTypes, optional): type of basis. Defaults to None, in which case the fock basis is used.
    """

    if N_pre_diag is None:
        N_pre_diag = N

    assert N_pre_diag >= N, "N_pre_diag must be greater than or equal to N."

    _basis = basis if basis is not None else cls.DEFAULT_BASIS
    _hamiltonian = (
        hamiltonian if hamiltonian is not None else cls.DEFAULT_HAMILTONIAN
    )

    if use_linear:
        _hamiltonian = HamiltonianTypes.linear

    cls.param_validation(N, N_pre_diag, params, _hamiltonian, _basis)

    return cls(N, N_pre_diag, params, label, _basis, _hamiltonian)

get_H()

Return the Hamiltonian truncated in its eigenbasis.

Source code in jaxquantum/devices/base/base.py

def get_H(self):
    """
    Return the Hamiltonian truncated in its eigenbasis.
    """
    eig_systems = self.eig_systems
    values = eig_systems["vals"][..., : self.N]
    dtype = eig_systems["vecs"].dtype
    data = values.astype(dtype)[..., :, None] * jnp.eye(self.N, dtype=dtype)
    qdims = Qdims(((self.N,), (self.N,)))
    return Qarray._from_impl(DenseImpl._make(data), qdims)

get_H_full() abstractmethod

Return full H.

Source code in jaxquantum/devices/base/base.py

@abstractmethod
def get_H_full(self):
    """Return full H."""

get_H_linear() abstractmethod

Return linear terms in H.

Source code in jaxquantum/devices/base/base.py

@abstractmethod
def get_H_linear(self):
    """Return linear terms in H."""

get_linear_frequency() abstractmethod

Get frequency of linear terms.

Source code in jaxquantum/devices/base/base.py

@abstractmethod
def get_linear_frequency(self):
    """Get frequency of linear terms."""

param_validation(N, N_pre_diag, params, hamiltonian, basis) classmethod

This can be overridden by subclasses.

Source code in jaxquantum/devices/base/base.py

@classmethod
def param_validation(cls, N, N_pre_diag, params, hamiltonian, basis):
    """This can be overridden by subclasses."""
    pass

Drive

Bases: ABC

Source code in jaxquantum/devices/superconducting/drive.py

@struct.dataclass
class Drive(ABC):
    N: int = struct.field(pytree_node=False)
    fd: float
    _label: int = struct.field(pytree_node=False)

    @classmethod
    def create(cls, M_max, fd, label=0):
        cls.M_max = M_max
        N = 2 * M_max + 1
        return cls(N, fd, label)

    @property
    def label(self):
        return self.__class__.__name__ + str(self._label)

    @property
    def ops(self):
        return self.common_ops()

    def common_ops(self) -> Dict[str, Qarray]:
        ops = {}

        M_max = self.M_max

        # Construct M = ∑ₘ m|m><m| operator in drive charge basis
        ops["M"] = jnp2jqt(jnp.diag(jnp.arange(-M_max, M_max + 1)))

        # Construct Id = ∑ₘ|m><m| in the drive charge basis
        ops["id"] = jnp2jqt(jnp.identity(2 * M_max + 1))

        # Construct M₊ ≡ exp(iθ) and M₋ ≡ exp(-iθ) operators for drive
        ops["M-"] = jnp2jqt(jnp.eye(2 * M_max + 1, k=1))
        ops["M+"] = jnp2jqt(jnp.eye(2 * M_max + 1, k=-1))

        # Construct cos(θ) ≡ 1/2 * [M₊ + M₋] = 1/2 * ∑ₘ|m+1><m| + h.c
        ops["cos(θ)"] = 0.5 * (ops["M+"] + ops["M-"])

        # Construct sin(θ) ≡ -i/2 * [M₊ - M₋] = -i/2 * ∑ₘ|m+1><m| + h.c
        ops["sin(θ)"] = -0.5j * (ops["M+"] - ops["M-"])

        # Construct more general drive operators cos(kθ) and sin(kθ)
        for k in range(2, M_max + 1):
            ops[f"M_+{k}"] = jnp2jqt(jnp.eye(2 * M_max + 1, k=-k))
            ops[f"M_-{k}"] = jnp2jqt(jnp.eye(2 * M_max + 1, k=k))
            ops[f"cos({k}θ)"] = 0.5 * (ops[f"M_+{k}"] + ops[f"M_-{k}"])
            ops[f"sin({k}θ)"] = -0.5j * (ops[f"M_+{k}"] - ops[f"M_-{k}"])

        return ops

    #############################################################

    def get_H(self):
        """
        Bare "drive" Hamiltonian (fd * M) in the extended Hilbert space.
        """
        return self.fd * self.ops["M"]

get_H()

Bare "drive" Hamiltonian (fd * M) in the extended Hilbert space.

Source code in jaxquantum/devices/superconducting/drive.py

def get_H(self):
    """
    Bare "drive" Hamiltonian (fd * M) in the extended Hilbert space.
    """
    return self.fd * self.ops["M"]

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)

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)

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)))