feat(typing): add vdot to BaseFakeNumpyNamespace

Author: alexfikl

arraycontext/fake_numpy.py

@@ -312,6 +312,11 @@ def inner(a: ArrayOrScalar) -> ArrayOrScalar:
         # subclasses. Defining them as abstract methods would define them
         # as attributes, making __getattr__ fail to retrieve the intended function.
 
+        def vdot(self,
+                a: ArrayOrContainerOrScalarT,
+                b: ArrayOrContainerOrScalarT, /
+            ) -> ArrayOrScalar: ...
+
         def broadcast_to(self,
                 array: ArrayOrContainerOrScalar, /,
                 shape: tuple[int, ...]

arraycontext/impl/jax/fake_numpy.py

@@ -140,7 +140,8 @@ def stack(self, arrays, axis=0):
 
     # {{{ linear algebra
 
-    def vdot(self, x, y, dtype=None):
+    @override
+    def vdot(self, a, b, dtype=None):
         from arraycontext import rec_multimap_reduce_array_container
 
         def _rec_vdot(ary1, ary2):
@@ -151,7 +152,7 @@ def _rec_vdot(ary1, ary2):
 
             return jnp.vdot(ary1, ary2)
 
-        return rec_multimap_reduce_array_container(sum, _rec_vdot, x, y)
+        return rec_multimap_reduce_array_container(sum, _rec_vdot, a, b)
 
     # }}}
 

arraycontext/impl/numpy/fake_numpy.py

@@ -233,8 +233,9 @@ def inner_ravel(ary: ArrayOrScalar) -> ArrayOrScalar:
 
         return rec_map_container(inner_ravel, a)
 
-    def vdot(self, x, y):
-        return rec_multimap_reduce_array_container(sum, np.vdot, x, y)
+    @override
+    def vdot(self, a, b):
+        return rec_multimap_reduce_array_container(sum, np.vdot, a, b)
 
     def any(self, a, /):
         return rec_map_reduce_array_container(partial(reduce, np.logical_or),

arraycontext/impl/pyopencl/fake_numpy.py

@@ -62,6 +62,7 @@
     from arraycontext.typing import (
         Array,
         ArrayOrContainerOrScalar,
+        ArrayOrContainerOrScalarT,
         ArrayOrScalar,
     )
 
@@ -187,11 +188,14 @@ def stack(self, arrays, axis=0):
 
     # {{{ linear algebra
 
-    def vdot(self, x, y, dtype=None):
+    @override
+    def vdot(self,
+             a: ArrayOrContainerOrScalarT, b: ArrayOrContainerOrScalarT,
+             dtype: DTypeLike | None = None):
         return rec_multimap_reduce_array_container(
                 sum,
                 partial(cl_array.vdot, dtype=dtype, queue=self._array_context.queue),
-                x, y)
+                a, b)
 
     # }}}
 

arraycontext/impl/pytato/fake_numpy.py

@@ -301,7 +301,8 @@ def min(self,
     def absolute(self, a):
         return self.abs(a)
 
-    def vdot(self, a: Array, b: Array):
+    @override
+    def vdot(self, a: ArrayOrContainerOrScalar, b: ArrayOrContainerOrScalar):
         return rec_multimap_array_container(pt.vdot, a, b)
 
     # }}}