library

Library imports

Gate

Source code in jaxquantum/circuits/gates.py

@struct.dataclass
class Gate:
    dims: List[int] = struct.field(pytree_node=False)
    _U: Optional[Array] # Unitary
    _Ht: Optional[Array] # Hamiltonian
    _KM: Optional[Qarray] # Kraus map
    _c_ops: Optional[Qarray]
    _params: Dict[str, Any]
    _ts: Array
    _name: str = struct.field(pytree_node=False)
    num_modes: int = struct.field(pytree_node=False)

    @classmethod
    def create(
        cls,
        dims: Union[int, List[int]],
        name: str = "Gate",
        params: Optional[Dict[str, Any]] = None,
        ts: Optional[Array] = None,
        gen_U: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
        gen_Ht: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
        gen_c_ops: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
        gen_KM: Optional[Callable[[Dict[str, Any]], List[Qarray]]] = None,
        num_modes: int = 1,
    ):
        """Create a gate.

        Args:
            dims: Dimensions of the gate.
            name: Name of the gate.
            params: Parameters of the gate.
            ts: Times of the gate.
            gen_U: Function to generate the unitary of the gate.
            gen_Ht: Function to generate a function Ht(t) that takes in a time t and outputs a Hamiltonian Qarray.
            gen_KM: Function to generate the Kraus map of the gate.
            num_modes: Number of modes of the gate.
        """

        # TODO: add params to device?

        if isinstance(dims, int):
            dims = [dims]

        assert len(dims) == num_modes, (
            "Number of dimensions must match number of modes."
        )

        # Unitary
        _U = gen_U(params) if gen_U is not None else None 
        _Ht = gen_Ht(params) if gen_Ht is not None else None 
        _c_ops = gen_c_ops(params) if gen_c_ops is not None else Qarray.from_list([])

        if gen_KM is not None:
            _KM = gen_KM(params)
        elif _U is not None:
            _KM = Qarray.from_list([_U])

        return Gate(
            dims = dims,
            _U = _U,
            _Ht = _Ht,
            _KM = _KM,
            _c_ops = _c_ops,
            _params = params if params is not None else {},
            _ts=ts if ts is not None else jnp.array([]),
            _name=name,
            num_modes=num_modes,
        )

    def __str__(self):
        return self._name

    def __repr__(self):
        return self._name

    @property
    def name(self):
        return self._name

    @property
    def U(self):
        return self._U

    @property
    def Ht(self):
        return self._Ht

    @property
    def KM(self):
        return self._KM

    @property
    def c_ops(self):
        return self._c_ops

    @property
    def params(self):
        return self._params

    @property
    def ts(self):
        return self._ts

    def add_Ht(self, Ht: Callable[[float], Qarray]):
        """Add a Hamiltonian function to the gate."""
        def new_Ht(t):
            return Ht(t) + self.Ht(t) if self.Ht is not None else Ht(t)

        return Gate(
            dims = self.dims,
            _U = self.U,
            _Ht = new_Ht,
            _KM = self.KM,
            _c_ops = self.c_ops,
            _params = self.params,
            _ts = self.ts,
            _name = self.name,
            num_modes = self.num_modes,
        )

    def add_c_ops(self, c_ops: Qarray):
        """Add a c_ops to the gate."""
        return Gate(
            dims = self.dims,
            _U = self.U,
            _Ht = self.Ht,
            _KM = self.KM,
            _c_ops = concatenate([self.c_ops, c_ops]),
            _params = self.params,
            _ts = self.ts,
            _name = self.name,
            num_modes = self.num_modes,
        )

    def copy(self):
        """Return a copy of the gate."""
        return Gate(
            dims = deepcopy(self.dims),
            _U = self.U,
            _Ht = deepcopy(self.Ht),
            _KM = self.KM,
            _c_ops = self.c_ops,
            _params = deepcopy(self.params),
            _ts = self.ts,
            _name = self.name,
            num_modes = self.num_modes,
        )

add_Ht(Ht)

Add a Hamiltonian function to the gate.

Source code in jaxquantum/circuits/gates.py

def add_Ht(self, Ht: Callable[[float], Qarray]):
    """Add a Hamiltonian function to the gate."""
    def new_Ht(t):
        return Ht(t) + self.Ht(t) if self.Ht is not None else Ht(t)

    return Gate(
        dims = self.dims,
        _U = self.U,
        _Ht = new_Ht,
        _KM = self.KM,
        _c_ops = self.c_ops,
        _params = self.params,
        _ts = self.ts,
        _name = self.name,
        num_modes = self.num_modes,
    )

add_c_ops(c_ops)

Add a c_ops to the gate.

Source code in jaxquantum/circuits/gates.py

def add_c_ops(self, c_ops: Qarray):
    """Add a c_ops to the gate."""
    return Gate(
        dims = self.dims,
        _U = self.U,
        _Ht = self.Ht,
        _KM = self.KM,
        _c_ops = concatenate([self.c_ops, c_ops]),
        _params = self.params,
        _ts = self.ts,
        _name = self.name,
        num_modes = self.num_modes,
    )

copy()

Return a copy of the gate.

Source code in jaxquantum/circuits/gates.py

def copy(self):
    """Return a copy of the gate."""
    return Gate(
        dims = deepcopy(self.dims),
        _U = self.U,
        _Ht = deepcopy(self.Ht),
        _KM = self.KM,
        _c_ops = self.c_ops,
        _params = deepcopy(self.params),
        _ts = self.ts,
        _name = self.name,
        num_modes = self.num_modes,
    )

create(dims, name='Gate', params=None, ts=None, gen_U=None, gen_Ht=None, gen_c_ops=None, gen_KM=None, num_modes=1) classmethod

Create a gate.

Parameters:

Name	Type	Description	Default
dims	Union[int, List[int]]	Dimensions of the gate.	required
name	str	Name of the gate.	'Gate'
params	Optional[Dict[str, Any]]	Parameters of the gate.	None
ts	Optional[Array]	Times of the gate.	None
gen_U	Optional[Callable[[Dict[str, Any]], Qarray]]	Function to generate the unitary of the gate.	None
gen_Ht	Optional[Callable[[Dict[str, Any]], Qarray]]	Function to generate a function Ht(t) that takes in a time t and outputs a Hamiltonian Qarray.	None
gen_KM	Optional[Callable[[Dict[str, Any]], List[Qarray]]]	Function to generate the Kraus map of the gate.	None
num_modes	int	Number of modes of the gate.	1

Source code in jaxquantum/circuits/gates.py

@classmethod
def create(
    cls,
    dims: Union[int, List[int]],
    name: str = "Gate",
    params: Optional[Dict[str, Any]] = None,
    ts: Optional[Array] = None,
    gen_U: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
    gen_Ht: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
    gen_c_ops: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
    gen_KM: Optional[Callable[[Dict[str, Any]], List[Qarray]]] = None,
    num_modes: int = 1,
):
    """Create a gate.

    Args:
        dims: Dimensions of the gate.
        name: Name of the gate.
        params: Parameters of the gate.
        ts: Times of the gate.
        gen_U: Function to generate the unitary of the gate.
        gen_Ht: Function to generate a function Ht(t) that takes in a time t and outputs a Hamiltonian Qarray.
        gen_KM: Function to generate the Kraus map of the gate.
        num_modes: Number of modes of the gate.
    """

    # TODO: add params to device?

    if isinstance(dims, int):
        dims = [dims]

    assert len(dims) == num_modes, (
        "Number of dimensions must match number of modes."
    )

    # Unitary
    _U = gen_U(params) if gen_U is not None else None 
    _Ht = gen_Ht(params) if gen_Ht is not None else None 
    _c_ops = gen_c_ops(params) if gen_c_ops is not None else Qarray.from_list([])

    if gen_KM is not None:
        _KM = gen_KM(params)
    elif _U is not None:
        _KM = Qarray.from_list([_U])

    return Gate(
        dims = dims,
        _U = _U,
        _Ht = _Ht,
        _KM = _KM,
        _c_ops = _c_ops,
        _params = params if params is not None else {},
        _ts=ts if ts is not None else jnp.array([]),
        _name=name,
        num_modes=num_modes,
    )

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
    @overload
    def create(cls, data, dims=None, bdims=None, implementation: Literal[QarrayImplType.DENSE] = QarrayImplType.DENSE) -> "Qarray[DenseImpl]":
        ...

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

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

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

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

        Args:
            data: Input data array (dense array-like or ``sparse.BCOO``).
            dims: Quantum dimensions as ``((row_dims...), (col_dims...))``.
                Inferred from *data* shape when ``None``.
            bdims: Tuple of batch dimension sizes.  Inferred from the leading
                dimensions of *data* when ``None``.
            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.
        """
        # Step 1: Prepare data ----
        data = robust_asarray(data)

        if len(data.shape) == 1 and data.shape[0] > 0:
            data = data.reshape(data.shape[0], 1)

        if len(data.shape) >= 2:
            if 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:
            if len(data.shape) - len(bdims) == 1:
                data = data.reshape(*data.shape[:-1], data.shape[-1], 1)
        # ----

        # Step 2: Prepare dimensions ----
        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)):
            # This handles the case where only the hilbert space dimensions are 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]))

        check_dims(dims, bdims, data.shape)

        qdims = Qdims(dims)

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

    @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]"]) -> "Qarray[DenseImpl]":
        ...

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

    @classmethod
    def from_list(cls, qarr_list: List[Qarray]) -> 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.

        Args:
            qarr_list: List of ``Qarray`` objects with identical ``dims`` and
                ``bdims``.  May be empty.

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

        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, 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, 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, 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:
            return Qarray.create(
                self.data[index],
                dims=self.dims,
                implementation=self.impl_type,
            )
        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
        else:
            new_shape = new_bdims + (prod(self.dims[0]),) + (-1,)

        # Preserve implementation type
        implementation = self.impl_type
        return Qarray.create(
            self.data.reshape(new_shape),
            dims=self.dims,
            bdims=new_bdims,
            implementation=implementation,
        )

    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 in [Qtypes.oper, 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)
        new_bdims = self.bdims

        # Preserve implementation type
        implementation = self.impl_type
        return Qarray.create(self.data, dims=new_qdims, bdims=new_bdims, implementation=implementation)

    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("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
        new_impl = self._impl.matmul(other._impl)

        return Qarray.create(
            new_impl.data,
            dims=_qdims_new.dims,
            implementation=new_impl.impl_type,
        )

    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
            other = other.reshape(other.shape + (1, 1))

        new_impl = self._impl.mul(other)
        return Qarray.create(
            new_impl.data,
            dims=self._qdims.dims,
            implementation=new_impl.impl_type,
        )

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

        return self.__mul__(1 / other)

    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.create(
                new_impl.data,
                dims=self.dims,
                implementation=new_impl.impl_type,
            )

        if robust_isscalar(other) and other == 0:
            return self.copy()

        if self.data.shape[-2] == self.data.shape[-1]:
            other = other + 0.0j
            if not robust_isscalar(other) and len(other.shape) > 0:  # not a scalar
                other = other.reshape(other.shape + (1, 1))
            eye_data = self._impl._eye_data(self.data.shape[-2], dtype=self.data.dtype)
            other = Qarray.create(
                other * eye_data,
                dims=self.dims,
                implementation=self.impl_type
            )
            return self.__add__(other)

        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.create(
                new_impl.data,
                dims=self.dims,
                implementation=new_impl.impl_type,
            )

        if robust_isscalar(other) and other == 0:
            return self.copy()

        if self.data.shape[-2] == self.data.shape[-1]:
            other = other + 0.0j

            if not robust_isscalar(other) and len(other.shape) > 0:  # not a scalar
                other = other.reshape(other.shape + (1, 1))
            eye_data = self._impl._eye_data(self.data.shape[-2], dtype=self.data.dtype)
            other = Qarray.create(
                other * eye_data,
                dims=self.dims,
                implementation=self.impl_type
            )
            return self.__sub__(other)

        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.create(
            new_impl.data,
            dims=self.dims,
            implementation=new_impl.impl_type,
        )

    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.create(
            new_impl.data,
            dims=self.dims,
            implementation=new_impl.impl_type,
        )

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

        Returns:
            A new ``Qarray`` containing the complex-conjugated elements.
        """
        new_impl = self._impl.conj()
        return Qarray.create(
            new_impl.data,
            dims=self.dims,
            implementation=new_impl.impl_type,
        )

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

    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.create(
        new_impl.data,
        dims=self.dims,
        implementation=new_impl.impl_type,
    )

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, implementation=QarrayImplType.DENSE) classmethod

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

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

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

Create a Qarray from raw data.

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

Parameters:

Name	Type	Description	Default
data		Input data array (dense array-like or sparse.BCOO).	required
dims		Quantum dimensions as ((row_dims...), (col_dims...)). Inferred from data shape when None.	None
bdims		Tuple of batch dimension sizes. Inferred from the leading dimensions of data when None.	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, implementation=QarrayImplType.DENSE):
    """Create a ``Qarray`` from raw data.

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

    Args:
        data: Input data array (dense array-like or ``sparse.BCOO``).
        dims: Quantum dimensions as ``((row_dims...), (col_dims...))``.
            Inferred from *data* shape when ``None``.
        bdims: Tuple of batch dimension sizes.  Inferred from the leading
            dimensions of *data* when ``None``.
        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.
    """
    # Step 1: Prepare data ----
    data = robust_asarray(data)

    if len(data.shape) == 1 and data.shape[0] > 0:
        data = data.reshape(data.shape[0], 1)

    if len(data.shape) >= 2:
        if 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:
        if len(data.shape) - len(bdims) == 1:
            data = data.reshape(*data.shape[:-1], data.shape[-1], 1)
    # ----

    # Step 2: Prepare dimensions ----
    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)):
        # This handles the case where only the hilbert space dimensions are 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]))

    check_dims(dims, bdims, data.shape)

    qdims = Qdims(dims)

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

from_list(qarr_list: List['Qarray[DenseImpl]']) -> 'Qarray[DenseImpl]'

from_list(qarr_list: List['Qarray[SparseBCOOImpl]']) -> '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.

Parameters:

Name	Type	Description	Default
qarr_list	List[Qarray]	List of Qarray objects with identical dims and bdims. May be empty.	required

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

    Args:
        qarr_list: List of ``Qarray`` objects with identical ``dims`` and
            ``bdims``.  May be empty.

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

    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, 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, 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, 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.create(
        new_impl.data,
        dims=self.dims,
        implementation=new_impl.impl_type,
    )

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.create(
        new_impl.data,
        dims=self.dims,
        implementation=new_impl.impl_type,
    )

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
    else:
        new_shape = new_bdims + (prod(self.dims[0]),) + (-1,)

    # Preserve implementation type
    implementation = self.impl_type
    return Qarray.create(
        self.data.reshape(new_shape),
        dims=self.dims,
        bdims=new_bdims,
        implementation=implementation,
    )

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)
    new_bdims = self.bdims

    # Preserve implementation type
    implementation = self.impl_type
    return Qarray.create(self.data, dims=new_qdims, bdims=new_bdims, implementation=implementation)

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 in [Qtypes.oper, 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)

CD(N, beta, ts=None)

Conditional displacement gate.

Parameters:

Name	Type	Description	Default
N		Hilbert space dimension.	required
beta		Conditional displacement amplitude.	required
ts		Optional time sequence for hamiltonian simulation.	None

Returns:

Type	Description
	Conditional displacement gate.

Source code in jaxquantum/circuits/library/oscillator.py

def CD(N, beta, ts=None):
    """Conditional displacement gate.

    Args:
        N: Hilbert space dimension.
        beta: Conditional displacement amplitude.
        ts: Optional time sequence for hamiltonian simulation.

    Returns:
        Conditional displacement gate.
    """
    g = basis(2, 0)
    e = basis(2, 1)

    gg = g @ g.dag()
    ee = e @ e.dag()

    gen_Ht = None
    if ts is not None:
        delta_t = ts[-1] - ts[0]
        amp = 1j * beta / delta_t / 2
        a = destroy(N)
        gen_Ht = lambda params: lambda t: (
            gg
            ^ (jnp.conj(amp) * a + amp * a.dag()) + ee
            ^ (jnp.conj(-amp) * a + (-amp) * a.dag())
        )

    return Gate.create(
        [2, N],
        name="CD",
        params={"beta": beta},
        gen_U=lambda params: (gg ^ displace(N, params["beta"] / 2))
        + (ee ^ displace(N, -params["beta"] / 2)),
        gen_Ht=gen_Ht,
        ts=ts,
        num_modes=2,
    )

CR(N, theta)

Conditional rotation gate.

Parameters:

Name	Type	Description	Default
N		Hilbert space dimension.	required
theta		Conditional rotation angle.	required

Returns:

Type	Description
	Conditional rotation gate.

Source code in jaxquantum/circuits/library/oscillator.py

def CR(N, theta):
    """Conditional rotation gate.

    Args:
        N: Hilbert space dimension.
        theta: Conditional rotation angle.

    Returns:
        Conditional rotation gate.
    """
    g = basis(2, 0)
    e = basis(2, 1)

    gg = g @ g.dag()
    ee = e @ e.dag()


    return Gate.create(
        [2, N],
        name="CR",
        params={"theta": theta},
        gen_U=lambda params: (gg ^ (-1.j*theta/2*create(N)@destroy(N)).expm())
        + (ee ^ (1.j*theta/2*create(N)@destroy(N)).expm()),
        num_modes=2,
    )

D(N, alpha, ts=None, c_ops=None)

Displacement gate.

Parameters:

Name	Type	Description	Default
N		Hilbert space dimension.	required
alpha		Displacement amplitude.	required
ts		Optional time array for hamiltonian simulation.	None
c_ops		Optional collapse operators.	None

Returns:

Type	Description
	Displacement gate.

Source code in jaxquantum/circuits/library/oscillator.py

def D(N, alpha, ts=None, c_ops=None):
    """Displacement gate.

    Args:
        N: Hilbert space dimension.
        alpha: Displacement amplitude.
        ts: Optional time array for hamiltonian simulation.
        c_ops: Optional collapse operators.

    Returns:
        Displacement gate.
    """
    gen_Ht = None
    if ts is not None:
        delta_t = ts[-1] - ts[0]
        amp = 1j * alpha / delta_t
        a = destroy(N)
        gen_Ht = lambda params: (lambda t: jnp.conj(amp) * a + amp * a.dag())

    return Gate.create(
        N,
        name="D",
        params={"alpha": alpha},
        gen_U=lambda params: displace(N, params["alpha"]),
        gen_Ht=gen_Ht,
        ts=ts,
        gen_c_ops=lambda params: Qarray.from_list([]) if c_ops is None else c_ops,
        num_modes=1,
    )

ECD(N, beta, ts=None)

Echoed conditional displacement gate.

Parameters:

Name	Type	Description	Default
N		Hilbert space dimension.	required
beta		Conditional displacement amplitude.	required
ts		Optional time sequence for hamiltonian simulation.	None

Returns:

Type	Description
	Echoed conditional displacement gate.

Source code in jaxquantum/circuits/library/oscillator.py

def ECD(N, beta, ts=None):
    """Echoed conditional displacement gate.

    Args:
        N: Hilbert space dimension.
        beta: Conditional displacement amplitude.
        ts: Optional time sequence for hamiltonian simulation.

    Returns:
        Echoed conditional displacement gate.
    """
    g = basis(2, 0)
    e = basis(2, 1)

    eg = e @ g.dag()
    ge = g @ e.dag()

    # gen_Ht = None
    # if ts is not None:
    #     delta_t = ts[-1] - ts[0]
    #     amp = 1j * beta / delta_t / 2
    #     a = destroy(N)
    #     gen_Ht = lambda params: lambda t: (
    #         eg
    #         ^ (jnp.conj(amp) * a + amp * a.dag()) + ge
    #         ^ (jnp.conj(-amp) * a + (-amp) * a.dag())
    #     )

    return Gate.create(
        [2, N],
        name="ECD",
        params={"beta": beta},
        gen_U=lambda params: (eg ^ displace(N, params["beta"] / 2))
        + (ge ^ displace(N, -params["beta"] / 2)),
        gen_Ht=None,
        ts=ts,
        num_modes=2,
    )

basis(N, k, implementation=QarrayImplType.DENSE)

Creates a |k> (i.e. fock state) ket in a specified Hilbert Space.

Parameters:

Name	Type	Description	Default
N	int	Hilbert space dimension	required
k	int	fock number	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
	Fock State |k>

Source code in jaxquantum/core/operators.py

def basis(N: int, k: int, implementation: QarrayImplType = QarrayImplType.DENSE):
    """Creates a |k> (i.e. fock state) ket in a specified Hilbert Space.

    Args:
        N: Hilbert space dimension
        k: fock number
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        Fock State |k>
    """
    return Qarray.create(one_hot(k, N).reshape(N, 1), implementation=implementation)

create(N, implementation=QarrayImplType.DENSE)

creation operator

Parameters:

Name	Type	Description	Default
N		Hilbert space size	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	creation operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def create(N, implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """creation operator

    Args:
        N: Hilbert space size
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        creation operator in Hilber Space of size N
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # Single subdiagonal at offset -1; Convention A: 1 trailing zero.
        diags = jnp.zeros((1, N), dtype=jnp.float64)
        diags = diags.at[0, :N - 1].set(jnp.sqrt(jnp.arange(1, N, dtype=jnp.float64)))
        return _make_sparsedia(offsets=(-1,), diags=diags)
    return Qarray.create(jnp.diag(jnp.sqrt(jnp.arange(1, N)), k=-1), implementation=implementation)

destroy(N, implementation=QarrayImplType.DENSE)

annihilation operator

Parameters:

Name	Type	Description	Default
N		Hilbert space size	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	annilation operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def destroy(N, implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """annihilation operator

    Args:
        N: Hilbert space size
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        annilation operator in Hilber Space of size N
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # Single superdiagonal at offset +1; Convention A: 1 leading zero.
        diags = jnp.zeros((1, N), dtype=jnp.float64)
        diags = diags.at[0, 1:].set(jnp.sqrt(jnp.arange(1, N, dtype=jnp.float64)))
        return _make_sparsedia(offsets=(1,), diags=diags)
    return Qarray.create(jnp.diag(jnp.sqrt(jnp.arange(1, N)), k=1), implementation=implementation)

diag_expm(diag_matrix)

Computes expm of a diagonal matrix efficiently (O(N) instead of O(N^3)).

Source code in jaxquantum/circuits/library/oscillator.py

def diag_expm(diag_matrix):
    """Computes expm of a diagonal matrix efficiently (O(N) instead of O(N^3))."""
    # Extract diagonal, exponentiate elements, put back on diagonal
    return jnp.diag(jnp.exp(jnp.diagonal(diag_matrix)))

displace(N, α)

Displacement operator

Parameters:

Name	Type	Description	Default
N		Hilbert Space Size	required
α		Phase space displacement	required

Returns:

Type	Description
Qarray	Displace operator D(α)

Source code in jaxquantum/core/operators.py

def displace(N, α) -> Qarray:
    """Displacement operator

    Args:
        N: Hilbert Space Size
        α: Phase space displacement

    Returns:
        Displace operator D(α)
    """
    a = destroy(N)
    return (α * a.dag() - jnp.conj(α) * a).expm()

hadamard(implementation=QarrayImplType.DENSE)

H

Returns:

Name	Type	Description
H	Qarray	Hadamard gate

Source code in jaxquantum/core/operators.py

def hadamard(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """H

    Returns:
        H: Hadamard gate
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        s = 1.0 / jnp.sqrt(2.0)
        # offset -1: valid at [0]   → diag[0]=A[1,0]=s, diag[1]=0 (trailing zero)
        # offset  0: valid at [0:2] → diag[0]=A[0,0]=s, diag[1]=A[1,1]=-s
        # offset +1: valid at [1]   → diag[0]=0 (leading zero), diag[1]=A[0,1]=s
        diags = jnp.array([[s, 0.0], [s, -s], [0.0, s]])
        return _make_sparsedia(offsets=(-1, 0, 1), diags=diags)
    return Qarray.create(jnp.array([[1, 1], [1, -1]]) / jnp.sqrt(2), implementation=implementation)

identity(*args, implementation=QarrayImplType.DENSE, **kwargs)

Identity matrix.

Parameters:

Name	Type	Description	Default
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	Identity matrix.

Source code in jaxquantum/core/operators.py

def identity(*args, implementation: QarrayImplType = QarrayImplType.DENSE, **kwargs) -> Qarray:
    """Identity matrix.

    Args:
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        Identity matrix.
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # jnp.eye(*args) is typically eye(N) or eye(N, N); extract N from args.
        n = args[0] if args else kwargs.get("N", kwargs.get("n", None))
        if n is not None and (len(args) <= 1) and not kwargs:
            diags = jnp.ones((1, int(n)), dtype=jnp.float64)
            return _make_sparsedia(offsets=(0,), diags=diags)
    return Qarray.create(jnp.eye(*args, **kwargs), implementation=implementation)

num(N, implementation=QarrayImplType.DENSE)

Number operator

Parameters:

Name	Type	Description	Default
N		Hilbert Space size	required
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	number operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def num(N, implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """Number operator

    Args:
        N: Hilbert Space size
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        number operator in Hilber Space of size N
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # Main diagonal only; no leading/trailing zeros needed (offset 0).
        diags = jnp.arange(N, dtype=jnp.float64).reshape(1, N)
        return _make_sparsedia(offsets=(0,), diags=diags)
    return Qarray.create(jnp.diag(jnp.arange(N)), implementation=implementation)

qubit_rotation(theta, nx, ny, nz)

Single qubit rotation.

Parameters:

Name	Type	Description	Default
theta	float	rotation angle.	required
nx		rotation axis x component.	required
ny		rotation axis y component.	required
nz		rotation axis z component.	required

Returns:

Type	Description
Qarray	Single qubit rotation operator.

Source code in jaxquantum/core/operators.py

def qubit_rotation(theta: float, nx, ny, nz) -> Qarray:
    """Single qubit rotation.

    Args:
        theta: rotation angle.
        nx: rotation axis x component.
        ny: rotation axis y component.
        nz: rotation axis z component.

    Returns:
        Single qubit rotation operator.
    """
    return jnp.cos(theta / 2) * identity(2) - 1j * jnp.sin(theta / 2) * (
        nx * sigmax() + ny * sigmay() + nz * sigmaz()
    )

sigmax(implementation=QarrayImplType.DENSE)

σx

Parameters:

Name	Type	Description	Default
implementation	QarrayImplType	Qarray implementation type, e.g. "sparse" or "dense".	DENSE

Returns:

Type	Description
Qarray	σx Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmax(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """σx

    Args:
        implementation: Qarray implementation type, e.g. "sparse" or "dense".

    Returns:
        σx Pauli Operator
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        # Offset -1: valid at [0:1] → diag[0] = A[1,0] = 1.0, diag[1] = 0 (trailing zero)
        # Offset +1: valid at [1:]  → diag[0] = 0 (leading zero), diag[1] = A[0,1] = 1.0
        diags = jnp.array([[1.0, 0.0], [0.0, 1.0]])
        return _make_sparsedia(offsets=(-1, 1), diags=diags)
    return Qarray.create(jnp.array([[0.0, 1.0], [1.0, 0.0]]), implementation=implementation)

sigmay(implementation=QarrayImplType.DENSE)

σy

Returns:

Type	Description
Qarray	σy Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmay(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """σy

    Returns:
        σy Pauli Operator
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        diags = jnp.array([[1.0j, 0.0], [0.0, -1.0j]])
        return _make_sparsedia(offsets=(-1, 1), diags=diags)
    return Qarray.create(jnp.array([[0.0, -1.0j], [1.0j, 0.0]]), implementation=implementation)

sigmaz(implementation=QarrayImplType.DENSE)

σz

Returns:

Type	Description
Qarray	σz Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmaz(implementation: QarrayImplType = QarrayImplType.DENSE) -> Qarray:
    """σz

    Returns:
        σz Pauli Operator
    """
    if QarrayImplType(implementation) == QarrayImplType.SPARSE_DIA:
        diags = jnp.array([[1.0, -1.0]])
        return _make_sparsedia(offsets=(0,), diags=diags)
    return Qarray.create(jnp.array([[1.0, 0.0], [0.0, -1.0]]), implementation=implementation)

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]
        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
            if len(a.shape) > len(b.shape):
                batch_dim = a.shape[:-2]
            elif len(a.shape) == len(b.shape):
                batch_dim = a.shape[:-2] if prod(a.shape[:-2]) > prod(b.shape[:-2]) else b.shape[:-2]
            else:
                batch_dim = b.shape[:-2]

            # NOTE: implementation einsum should be used when available
            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]
        return Qarray.create(data, dims=(dims_0, dims_1), implementation=target_impl_type)

    # 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.create(current_impl.data, dims=(dims_0, dims_1),
                         implementation=current_impl.impl_type)