Fix tests that fail if enable_checks is true under NumPy 2.0.0rc1.

np.vecdot is missing `__module__` under NumPy 2.0.0rc1. PiperOrigin-RevId: 621532796
jax-ml · Apr 3, 2024 · e2f4774 · e2f4774
1 parent d89f0d6
commit e2f4774
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
@@ -3443,8 +3443,11 @@ def vdot(
              preferred_element_type=preferred_element_type)
 
 
-@util.implements(getattr(np, "vecdot", None), lax_description=_PRECISION_DOC,
-                 extra_params=_DOT_PREFERRED_ELEMENT_TYPE_DESCRIPTION)
+@util.implements(
+    getattr(np, "vecdot", None), lax_description=_PRECISION_DOC,
+    extra_params=_DOT_PREFERRED_ELEMENT_TYPE_DESCRIPTION,
+    # TODO(phawkins): numpy.vecdot doesn't have a __module__ attribute.
+    module="numpy")
 def vecdot(x1: ArrayLike, x2: ArrayLike, /, *, axis: int = -1,
            precision: PrecisionLike = None,
            preferred_element_type: DTypeLike | None = None) -> Array: