Batch Refactor

mrmustard/lab_dev/circuit_components.py

@@ -336,7 +336,7 @@ def quadrature(self, quad: Batch[Vector], phi: float = 0.0) -> ComplexTensor:
             # Find where all the bras and kets are so they can be conjugated appropriately
             conjugates = [i not in self.wires.ket.indices for i in range(len(self.wires.indices))]
             quad_basis = math.sum(
-                [quadrature_basis(array, quad, conjugates, phi) for array in fock_arrays], axes=[0]
+                [quadrature_basis(array, quad, conjugates, phi) for array in fock_arrays], axis=[0]
             )
             return quad_basis
 

mrmustard/lab_dev/circuit_components_utils/trace_out.py

@@ -86,7 +86,7 @@ def __custom_rrshift__(self, other: CircuitComponent | complex) -> CircuitCompon
             ansatz = other.ansatz
             wires = other.wires
         elif not ket or not bra:
-            ansatz = other.ansatz.conj[idx_z] @ other.ansatz[idx_z]
+            ansatz = other.ansatz.conj.contract(other.ansatz, idx_z, idx_z)
             wires, _ = (other.wires.adjoint @ other.wires)[0] @ self.wires
         else:
             ansatz = other.ansatz.trace(idx_z, idx_zconj)

mrmustard/lab_dev/states/base.py

@@ -204,7 +204,6 @@ def from_fock(
         modes: Sequence[int],
         array: ComplexTensor,
         name: str | None = None,
-        batched: bool = False,
     ) -> State:
         r"""
         Initializes a state of type ``cls`` from an array parametrizing the
@@ -218,17 +217,16 @@ def from_fock(
 
             >>> modes = [0]
             >>> array = Coherent(modes, x=0.1).to_fock().ansatz.array
-            >>> coh = Ket.from_fock(modes, array, batched=True)
+            >>> coh = Ket.from_fock(modes, array)
 
             >>> assert coh.modes == modes
-            >>> assert coh.ansatz == ArrayAnsatz(array, batched=True)
+            >>> assert coh.ansatz == ArrayAnsatz(array)
             >>> assert isinstance(coh, Ket)
 
         Args:
             modes: The modes of this state.
             array: The Fock array.
             name: The name of this state.
-            batched: Whether the given array is batched.
 
         Returns:
             A state.
@@ -385,7 +383,7 @@ def visualize_2d(
         shape = [max(min_shape, d) for d in self.auto_shape()]
         state = self.to_fock(tuple(shape))
         state = state.dm()
-        dm = math.sum(state.ansatz.array, axes=[0])
+        dm = math.sum(state.ansatz.array, axis=[0])
 
         x, prob_x = quadrature_distribution(dm)
         p, prob_p = quadrature_distribution(dm, np.pi / 2)
@@ -501,7 +499,7 @@ def visualize_3d(
         shape = [max(min_shape, d) for d in self.auto_shape()]
         state = self.to_fock(tuple(shape))
         state = state.dm()
-        dm = math.sum(state.ansatz.array, axes=[0])
+        dm = math.sum(state.ansatz.array, axis=[0])
 
         xvec = np.linspace(*xbounds, resolution)
         pvec = np.linspace(*pbounds, resolution)
@@ -575,7 +573,7 @@ def visualize_dm(
             raise ValueError("DM visualization not available for multi-mode states.")
         state = self.to_fock(cutoff)
         state = state.dm()
-        dm = math.sum(state.ansatz.array, axes=[0])
+        dm = math.sum(state.ansatz.array, axis=[0])
 
         fig = go.Figure(
             data=go.Heatmap(z=abs(dm), colorscale="viridis", name="abs(ρ)", showscale=False)

mrmustard/lab_dev/states/dm.py

@@ -28,6 +28,7 @@
 from mrmustard.math.lattice.strategies.vanilla import autoshape_numba
 from mrmustard.physics.ansatz import ArrayAnsatz, PolyExpAnsatz
 from mrmustard.physics.bargmann_utils import wigner_to_bargmann_rho
+from mrmustard.physics.batches import Batch
 from mrmustard.physics.gaussian_integrals import complex_gaussian_integral_2
 from mrmustard.physics.representations import Representation
 from mrmustard.physics.wires import Wires
@@ -121,11 +122,10 @@ def from_bargmann(
     def from_fock(
         cls,
         modes: Sequence[int],
-        array: ComplexTensor,
+        array: ComplexTensor | Batch[ComplexTensor],
         name: str | None = None,
-        batched: bool = False,
     ) -> State:
-        return DM.from_ansatz(modes, ArrayAnsatz(array, batched), name)
+        return DM.from_ansatz(modes, ArrayAnsatz(array), name)
 
     @classmethod
     def from_ansatz(

mrmustard/lab_dev/states/ket.py

@@ -28,6 +28,7 @@
 from mrmustard.math.lattice.strategies.vanilla import autoshape_numba
 from mrmustard.physics.ansatz import ArrayAnsatz, PolyExpAnsatz
 from mrmustard.physics.bargmann_utils import wigner_to_bargmann_psi
+from mrmustard.physics.batches import Batch
 from mrmustard.physics.gaussian import purity
 from mrmustard.physics.representations import Representation
 from mrmustard.physics.wires import Wires
@@ -37,7 +38,6 @@
     ComplexTensor,
     RealVector,
     Scalar,
-    Batch,
 )
 
 from .base import State, _validate_operator, OperatorType
@@ -100,11 +100,10 @@ def from_bargmann(
     def from_fock(
         cls,
         modes: Sequence[int],
-        array: ComplexTensor,
+        array: ComplexTensor | Batch[ComplexTensor],
         name: str | None = None,
-        batched: bool = False,
     ) -> State:
-        return Ket.from_ansatz(modes, ArrayAnsatz(array, batched), name)
+        return Ket.from_ansatz(modes, ArrayAnsatz(array), name)
 
     @classmethod
     def from_ansatz(

mrmustard/lab_dev/transformations/base.py

@@ -201,7 +201,7 @@ def symplectic(self):
         Returns the symplectic matrix that corresponds to this unitary
         """
         batch_size = self.ansatz.batch_size
-        return [au2Symplectic(self.ansatz.A[batch, :, :]) for batch in range(batch_size)]
+        return [au2Symplectic(self.ansatz.A[batch]) for batch in range(batch_size)]
 
     @classmethod
     def from_bargmann(
@@ -380,7 +380,7 @@ def is_CP(self) -> bool:
             raise ValueError(
                 "Physicality conditions are not implemented for batch dimension larger than 1."
             )
-        A = self.ansatz.A
+        A = self.ansatz.A.data
         m = A.shape[-1] // 2
         gamma_A = A[0, :m, m:]
 
@@ -396,7 +396,7 @@ def is_TP(self) -> bool:
         r"""
         Whether this channel is trace preserving (TP).
         """
-        A = self.ansatz.A
+        A = self.ansatz.A.data
         m = A.shape[-1] // 2
         gamma_A = A[0, :m, m:]
         lambda_A = A[0, m:, m:]
@@ -521,7 +521,7 @@ def random(cls, modes: Sequence[int], max_r: float = 1.0) -> Channel:
         U = Unitary.random(range(3 * m), max_r)
         u_psi = Vacuum(range(2 * m)) >> U
         A = u_psi.ansatz
-        kraus = A.conj[range(2 * m)] @ A[range(2 * m)]
+        kraus = A.conj.contract(A, range(2 * m), range(2 * m))
         return Channel.from_bargmann(modes, modes, kraus.triple)
 
     def __rshift__(self, other: CircuitComponent) -> CircuitComponent:

mrmustard/lab_dev/transformations/dgate.py

@@ -24,8 +24,9 @@
 from mrmustard import math
 
 from .base import Unitary
-from ...physics.representations import Representation
 from ...physics.ansatz import PolyExpAnsatz, ArrayAnsatz
+from ...physics.batches import Batch
+from ...physics.representations import Representation
 from ...physics import triples, fock_utils
 from ..utils import make_parameter, reshape_params
 
@@ -103,7 +104,9 @@ def __init__(
             ),
         ).representation
 
-    def fock_array(self, shape: int | Sequence[int] = None, batched=False) -> ComplexTensor:
+    def fock_array(
+        self, shape: int | Sequence[int] = None, batched=False
+    ) -> ComplexTensor | Batch[ComplexTensor]:
         r"""
         Returns the unitary representation of the Displacement gate using the Laguerre polynomials.
         If the shape is not given, it defaults to the ``auto_shape`` of the component if it is
@@ -144,11 +147,11 @@ def fock_array(self, shape: int | Sequence[int] = None, batched=False) -> Comple
             )
         else:
             array = fock_utils.displacement(x[0], y[0], shape=shape)
-        arrays = math.expand_dims(array, 0) if batched else array
+        arrays = Batch(math.expand_dims(array, 0)) if batched else array
         return arrays
 
     def to_fock(self, shape: int | Sequence[int] | None = None) -> Dgate:
-        fock = ArrayAnsatz(self.fock_array(shape, batched=True), batched=True)
+        fock = ArrayAnsatz(self.fock_array(shape, batched=True))
         fock._original_abc_data = self.ansatz.triple
         ret = self._getitem_builtin(self.modes)
         ret._representation = Representation(fock, self.wires)

mrmustard/lab_dev/transformations/phasenoise.py

@@ -83,4 +83,4 @@ def __custom_rrshift__(self, other: CircuitComponent) -> CircuitComponent:
                 * self.phase_stdev.value**2
             )
             array *= phase_factors
-        return CircuitComponent(Representation(ArrayAnsatz(array, False), other.wires), self.name)
+        return CircuitComponent(Representation(ArrayAnsatz(array), other.wires), self.name)