nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: google/jax
Pull request #21856 (2024-06-13 17:46:57)

Deprecate passing of arrays in place of dtypes.
Related (if only tangentially) to #21824

In #21853 we made it so that non-hashable objects pass `check_valid_dtypes`, in order to treat traced and non-traced arrays the same.

But: passing an array in place of a dtype was only inadvertently supported! Because removing this support may break existing APIs, we need a deprecation cycle.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 536dba082..c9d48a974 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,8 @@ Remember to align the itemized text with the first line of an item within a list
     the last several JAX releases.
   * `jax.experimental.export` is deprecated. Use {mod}`jax.export` instead.
     See the [migration guide](https://jax.readthedocs.io/en/latest/export/export.html#migration-guide-from-jax-experimental-export).
+  * Passing an array in place of a dtype is now deprecated in most cases; e.g. for arrays
+    `x` and `y`, `x.astype(y)` will raise a warning. To silence it use `x.astype(y.dtype)`.
 
 ## jaxlib 0.4.30
 
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
index 0691d910b..1d50c5be7 100644
--- a/jax/_src/dtypes.py
+++ b/jax/_src/dtypes.py
@@ -31,7 +31,7 @@ import ml_dtypes
 import numpy as np
 
 from jax._src import config
-from jax._src.typing import DType, DTypeLike
+from jax._src.typing import Array, DType, DTypeLike
 from jax._src.util import set_module, StrictABC
 
 from jax._src import traceback_util
@@ -741,6 +741,12 @@ def result_type(*args: Any, return_weak_type_flag: bool = False) -> DType | tupl
   return (dtype, weak_type) if return_weak_type_flag else dtype  # type: ignore[return-value]
 
 def check_user_dtype_supported(dtype, fun_name=None):
+  if isinstance(dtype, Array):
+    # Deprecation warning added 2024 June 13.
+    warnings.warn("Passing an array as a dtype argument is deprecated; "
+                  "instead of dtype=arr use dtype=arr.dtype.",
+                  category=DeprecationWarning, stacklevel=3)
+    return  # no further check needed, as array dtypes have already been validated.
   if issubdtype(dtype, extended):
     return
   # Avoid using `dtype in [...]` because of numpy dtype equality overloading.


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
index 0a29e6168..4712e6aec 100644
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -565,6 +565,15 @@ class DtypesTest(jtu.JaxTestCase):
       dtype = np.dtype('float32')
     dtypes.check_user_dtype_supported(MyDtype())
 
+  def test_check_dtype_array(self):
+    x = jnp.arange(4)
+    msg = "Passing an array as a dtype argument is deprecated"
+    with self.assertWarnsRegex(DeprecationWarning, msg):
+      dtypes.check_user_dtype_supported(x)
+    with self.assertWarnsRegex(DeprecationWarning, msg):
+      jax.jit(dtypes.check_user_dtype_supported)(x)
+
+
 class EArrayTest(jtu.JaxTestCase):
 
   @parameterized.parameters([True, False])