Backend stack and broadcast to (#561)

Author: apchytr

mrmustard/math/backend_jax.py

@@ -87,6 +87,10 @@ def block(self, blocks: list[list[jnp.ndarray]], axes=(-2, -1)) -> jnp.ndarray:
         rows = [self.concat(row, axis=axes[1]) for row in blocks]
         return self.concat(rows, axis=axes[0])
 
+    @partial(jax.jit, static_argnames=["shape"])
+    def broadcast_to(self, array: jnp.ndarray, shape: tuple[int]) -> jnp.ndarray:
+        return jnp.broadcast_to(array, shape)
+
     @partial(jax.jit, static_argnames=["axis"])
     def prod(self, x: jnp.ndarray, axis: int | None):
         return jnp.prod(x, axis=axis)
@@ -417,6 +421,10 @@ def solve(self, matrix: jnp.ndarray, rhs: jnp.ndarray) -> jnp.ndarray:
     def sqrt(self, x: jnp.ndarray, dtype=None) -> jnp.ndarray:
         return jnp.sqrt(self.cast(x, dtype))
 
+    @partial(jax.jit, static_argnames=["axis"])
+    def stack(self, arrays: jnp.ndarray, axis: int = 0) -> jnp.ndarray:
+        return jnp.stack(arrays, axis=axis)
+
     @jax.jit
     def kron(self, tensor1: jnp.ndarray, tensor2: jnp.ndarray):
         return jnp.kron(tensor1, tensor2)

mrmustard/math/backend_manager.py

@@ -387,6 +387,18 @@ def block(self, blocks: list[list[Tensor]], axes=(-2, -1)) -> Tensor:
         """
         return self._apply("block", (blocks, axes))
 
+    def broadcast_to(self, array: Tensor, shape: tuple[int]) -> Tensor:
+        r"""Broadcasts an array to a new shape.
+
+        Args:
+            array: The array to broadcast.
+            shape: The shape to broadcast to.
+
+        Returns:
+            The broadcasted array.
+        """
+        return self._apply("broadcast_to", (array, shape))
+
     def cast(self, array: Tensor, dtype=None) -> Tensor:
         r"""Casts ``array`` to ``dtype``.
 
@@ -1145,17 +1157,18 @@ def sqrtm(self, tensor: Tensor, dtype=None) -> Tensor:
             The square root of ``x``"""
         return self._apply("sqrtm", (tensor, dtype))
 
-    def stack(self, values: Tensor, axis: int = 0) -> Tensor:
-        r"""Stacks a list of tensors along a new axis.
+    def stack(self, arrays: Sequence[Tensor], axis: int = 0) -> Tensor:
+        r"""Stack arrays in sequence along a new axis.
 
         Args:
-            values: A list of tensors to stack.
-            axis: The axis along which to introduce the new dimension.
+            arrays: Sequence of tensors to stack
+            axis: The axis along which to stack the arrays
 
         Returns:
-            The stacked tensor.
+            The stacked array
         """
-        return self._apply("stack", (values, axis))
+        arrays = self.astensor(arrays)
+        return self._apply("stack", (arrays, axis))
 
     def sum(self, array: Tensor, axis: int | Sequence[int] | None = None):
         r"""The sum of array.

mrmustard/math/backend_numpy.py

@@ -70,10 +70,6 @@ def abs(self, array: np.ndarray) -> np.ndarray:
         return np.abs(array)
 
     def allclose(self, array1: np.array, array2: np.array, atol: float, rtol: float) -> bool:
-        array1 = self.asnumpy(array1)
-        array2 = self.asnumpy(array2)
-        if array1.shape != array2.shape:
-            raise ValueError("Cannot compare arrays of different shapes.")
         return np.allclose(array1, array2, atol=atol, rtol=rtol)
 
     def any(self, array: np.ndarray) -> np.ndarray:
@@ -104,6 +100,9 @@ def block(self, blocks: list[list[np.ndarray]], axes=(-2, -1)) -> np.ndarray:
         rows = [self.concat(row, axis=axes[1]) for row in blocks]
         return self.concat(rows, axis=axes[0])
 
+    def broadcast_to(self, array: np.ndarray, shape: tuple[int]) -> np.ndarray:
+        return np.broadcast_to(array, shape)
+
     def block_diag(self, *blocks: list[np.ndarray]) -> np.ndarray:
         return sp.linalg.block_diag(*blocks)
 
@@ -400,8 +399,8 @@ def sort(self, array: np.ndarray, axis: int = -1) -> np.ndarray:
     def sqrt(self, x: np.ndarray, dtype=None) -> np.ndarray:
         return np.sqrt(self.cast(x, dtype))
 
-    def stack(self, values: list[np.ndarray], axis: int = 0) -> np.ndarray:
-        return np.stack(values, axis=axis)
+    def stack(self, arrays: np.ndarray, axis: int = 0) -> np.ndarray:
+        return np.stack(arrays, axis=axis)
 
     def sum(self, array: np.ndarray, axis: int | tuple[int] | None = None):
         return np.sum(array, axis=axis)

mrmustard/math/backend_tensorflow.py

@@ -71,10 +71,6 @@ def abs(self, array: tf.Tensor) -> tf.Tensor:
         return tf.abs(array)
 
     def allclose(self, array1: np.array, array2: np.array, atol: float, rtol: float) -> bool:
-        array1 = self.astensor(array1)
-        array2 = self.astensor(array2)
-        if array1.shape != array2.shape:
-            raise ValueError("Cannot compare arrays of different shapes.")
         return tf.experimental.numpy.allclose(array1, array2, atol=atol, rtol=rtol)
 
     def any(self, array: tf.Tensor) -> tf.Tensor:
@@ -110,6 +106,9 @@ def block(self, blocks: list[list[tf.Tensor]], axes=(-2, -1)) -> tf.Tensor:
         rows = [self.concat(row, axis=axes[1]) for row in blocks]
         return self.concat(rows, axis=axes[0])
 
+    def broadcast_to(self, array: tf.Tensor, shape: tuple[int]) -> tf.Tensor:
+        return tf.broadcast_to(array, shape)
+
     def boolean_mask(self, tensor: tf.Tensor, mask: tf.Tensor) -> Tensor:
         return tf.boolean_mask(tensor, mask)
 

tests/test_math/test_backend_manager.py

@@ -89,10 +89,10 @@ def test_allclose_error(self):
         arr1 = math.astensor([1, 2, 3])
         arr2 = math.astensor([[1, 2], [1, 2]])
 
-        if math.backend_name != "jax":
-            with pytest.raises(ValueError, match="Cannot compare"):
+        if math.backend_name == "numpy":
+            with pytest.raises(ValueError, match="could not be broadcast"):
                 math.allclose(arr1, arr2)
-        else:
+        elif math.backend_name == "jax":
             with pytest.raises(ValueError, match="Incompatible shapes"):
                 math.allclose(arr2, arr1)
 
@@ -262,6 +262,14 @@ def test_block_diag(self):
         assert R.shape == (8, 8)
         assert math.allclose(math.block([[I, O], [O, 1j * I]]), R)
 
+    def test_broadcast_to(self):
+        r"""
+        Tests the ``broadcast_to`` method.
+        """
+        arr = math.astensor([1, 2, 3])
+        res = math.broadcast_to(arr, (3, 3))
+        assert math.allclose(res, math.astensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]]))
+
     @pytest.mark.parametrize("t", types)
     def test_cast(self, t):
         r"""
@@ -674,6 +682,16 @@ def test_sqrtm(self):
         res = math.asnumpy(math.sqrtm(arr))
         assert math.allclose(res, 2 * np.eye(3))
 
+    def test_stack(self):
+        r"""
+        Tests the ``stack`` method.
+        """
+        arr1 = np.eye(3)
+        arr2 = 2 * np.eye(3)
+        res = math.asnumpy(math.stack([arr1, arr2], axis=0))
+        exp = np.stack([arr1, arr2], axis=0)
+        assert np.allclose(res, exp)
+
     def test_sum(self):
         r"""
         Tests the ``sum`` method.