Added Kronecker Product (#1728)

python/src/ops.cpp

@@ -1458,6 +1458,35 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The range of values.
       )pbdoc");
+  m.def(
+      "kron",
+      &kron,
+      nb::arg("a"),
+      nb::arg("b"),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def kron(a: array, b: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Compute the Kronecker product of two arrays `a` and `b`.
+        Args:
+          a (array): The first input array
+          b (array): The second input array
+          stream (Union[None, Stream, Device], optional): Optional stream or device for execution.
+            Default is `None`.
+        Returns:
+          array: The Kronecker product of `a` and `b`.
+        Examples:
+          >>> import mlx
+          >>> a = mlx.array([[1, 2], [3, 4]])
+          >>> b = mlx.array([[0, 5], [6, 7]])
+          >>> result = mlx.kron(a, b)
+          >>> print(result)
+          [[ 0  5  0 10]
+           [ 6  7 12 14]
+           [ 0 15  0 20]
+           [18 21 24 28]]
+      )pbdoc");
   m.def(
       "take",
       [](const mx::array& a,

python/tests/test_ops.py

@@ -1000,6 +1000,34 @@ class TestOps(mlx_tests.MLXTestCase):
 
         self.assertListEqual(mx.grad(func)(x).tolist(), expected)
 
+    def test_kron(self):
+        # Basic vector test
+        x = mx.array([1, 2])
+        y = mx.array([3, 4])
+        z = mx.kron(x, y)
+        self.assertEqual(z.tolist(), [3, 4, 6, 8])
+
+        # Basic matrix test
+        x = mx.array([[1, 2], [3, 4]])
+        y = mx.array([[0, 5], [6, 7]])
+        z = mx.kron(x, y)
+        self.assertEqual(
+            z.tolist(),
+            [[0, 5, 0, 10], [6, 7, 12, 14], [0, 15, 0, 20], [18, 21, 24, 28]],
+        )
+
+        # Test with different dimensions
+        x = mx.array([1, 2])  # (2,)
+        y = mx.array([[3, 4], [5, 6]])  # (2, 2)
+        z = mx.kron(x, y)
+        self.assertEqual(z.tolist(), [[3, 4, 6, 8], [5, 6, 10, 12]])
+
+        # Test with empty array
+        x = mx.array([])
+        y = mx.array([1, 2])
+        with self.assertRaises(ValueError):
+            mx.kron(x, y)
+
     def test_take(self):
         # Shape: 4 x 3 x 2
         l = [