Add vecdot to array API namespace (#3748)

Author: wgu9

docs/src/python/ops.rst

@@ -197,6 +197,7 @@ Operations
    triu
    unflatten
    unstack
+   vecdot
    var
    view
    where

mlx/ops.cpp

@@ -5534,6 +5534,28 @@ array inner(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   return tensordot(a, b, {-1}, {-1}, s);
 }
 
+array vecdot(
+    const array& a,
+    const array& b,
+    int axis /* = -1 */,
+    StreamOrDevice s /* = {} */) {
+  if (a.ndim() == 0 || b.ndim() == 0) {
+    throw std::invalid_argument("[vecdot] inputs must be at least 1D.");
+  }
+  int ax = axis < 0 ? axis + a.ndim() : axis;
+  if (ax < 0 || ax >= a.ndim()) {
+    throw std::invalid_argument("[vecdot] axis is out of bounds.");
+  }
+  if (axis < 0 ? axis + b.ndim() != ax : axis >= b.ndim()) {
+    throw std::invalid_argument("[vecdot] axis is out of bounds.");
+  }
+  if (a.shape(ax) != b.shape(ax)) {
+    throw std::invalid_argument(
+        "[vecdot] a and b must have the same size along axis.");
+  }
+  return sum(multiply(conjugate(a, s), b, s), ax, false, s);
+}
+
 /** Compute D = beta * C + alpha * (A @ B) */
 array addmm(
     array c,

mlx/ops.h

@@ -1574,6 +1574,10 @@ MLX_API array outer(const array& a, const array& b, StreamOrDevice s = {});
 /** Compute the inner product of two vectors. */
 MLX_API array inner(const array& a, const array& b, StreamOrDevice s = {});
 
+/** Compute a vector dot product along an axis. */
+MLX_API array
+vecdot(const array& a, const array& b, int axis = -1, StreamOrDevice s = {});
+
 /** Compute D = beta * C + alpha * (A @ B) */
 MLX_API array addmm(
     array c,

python/src/ops.cpp

@@ -4737,6 +4737,27 @@ void init_ops(nb::module_& m) {
       Returns:
         array: The inner product.
     )pbdoc");
+  m.def(
+      "vecdot",
+      &mx::vecdot,
+      nb::arg(),
+      nb::arg(),
+      "axis"_a = -1,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def vecdot(a: array, b: array, /, *, axis: int = -1, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+      Compute the vector dot product of two arrays along an axis.
+
+      Args:
+        a (array): Input array
+        b (array): Input array
+        axis (int, optional): Axis over which to compute the dot product. Default: ``-1``.
+
+      Returns:
+        array: The vector dot product.
+    )pbdoc");
   m.def(
       "outer",
       &mx::outer,

python/tests/test_array.py

@@ -192,7 +192,7 @@ def test_isdtype(self):
 
         # Reachable through the array API namespace.
         xp = mx.array(1.0).__array_namespace__()
-        for name in ("result_type", "can_cast", "isdtype"):
+        for name in ("result_type", "can_cast", "isdtype", "vecdot"):
             self.assertTrue(hasattr(xp, name), msg=name)
 
 

python/tests/test_ops.py

@@ -2655,6 +2655,27 @@ def test_inner(self):
         self.assertCmpNumpy([(1, 1, 2), (3, 2)], mx.inner, np.inner)
         self.assertCmpNumpy([(2, 3, 4), (4,)], mx.inner, np.inner)
 
+    def test_vecdot(self):
+        a = mx.array([[1, 2, 3], [4, 5, 6]])
+        b = mx.array([[7, 8, 9], [10, 11, 12]])
+        self.assertEqual(mx.vecdot(a, b).tolist(), [50, 167])
+        self.assertEqual(mx.vecdot(a, b, axis=0).tolist(), [47, 71, 99])
+
+        a = mx.array([1 + 2j, 3 + 4j])
+        b = mx.array([5 + 6j, 7 + 8j])
+        expected = np.vdot(np.array([1 + 2j, 3 + 4j]), np.array([5 + 6j, 7 + 8j]))
+        self.assertTrue(np.allclose(mx.vecdot(a, b), expected))
+
+        xp = mx.array(1.0).__array_namespace__()
+        self.assertEqual(xp.vecdot(mx.array([1, 2]), mx.array([3, 4])).item(), 11)
+
+        with self.assertRaises(ValueError):
+            mx.vecdot(mx.array(1), mx.array([1]))
+        with self.assertRaises(ValueError):
+            mx.vecdot(mx.array([1, 2]), mx.array([1, 2]), axis=1)
+        with self.assertRaises(ValueError):
+            mx.vecdot(mx.array([1, 2]), mx.array([1]))
+
     def test_outer(self):
         self.assertCmpNumpy([(3,), (3,)], mx.outer, np.outer)
         self.assertCmpNumpy(