Repository: jax-ml/jax (python)

Issue:
Tracers are now hashable.
### Description

`jax.jit(hash)(1)` does not error on any recent version of JAX. I went back and checked and the last time this seems to have produced an error was `0.4.8`.

Is this intended? It means that tracers do not duck-type as either concrete JAX arrays or as numpy arrays. In particular see #21825, which this caused.

### System info (python version, jaxlib version, accelerator, etc.)

jax:    0.4.29
jaxlib: 0.4.29
numpy:  1.26.4
python: 3.11.9 (main, Apr 19 2024, 11:43:47) [Clang 14.0.6 ]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Darwin', node='Air.localdomain', release='22.5.0', version='Darwin Kernel Version 22.5.0: Mon Apr 24 20:52:43 PDT 2023; root:xnu-8796.121.2~5/RELEASE_ARM64_T8112', machine='arm64')

Pull request description:
Make Tracers non-hashable
Hashing of tracers has been deprecated since JAX v0.4.30.

Fixes #21824

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8c34b8e36..bafe55110 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -36,6 +36,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
          use `uses_global_constants`.
       * the `lowering_platforms` kwarg for {func}`jax.export.export`: use
         `platforms` instead.
+  * Hashing of tracers, which has been deprecated since version 0.4.30, now
+    results in a `TypeError`.
 
 * New Features
   * {func}`jax.jit` got a new `compiler_options: dict[str, Any]` argument, for
diff --git a/jax/_src/basearray.py b/jax/_src/basearray.py
index c3145f32e..a89d4a294 100644
--- a/jax/_src/basearray.py
+++ b/jax/_src/basearray.py
@@ -53,6 +53,7 @@ class Array(abc.ABC):
   # associated basearray.pyi file.
 
   __slots__ = ['__weakref__']
+  __hash__ = None
 
   @property
   @abc.abstractmethod
diff --git a/jax/_src/core.py b/jax/_src/core.py
index cf570b8eb..6f96dc760 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -34,7 +34,6 @@ from weakref import ref
 
 import numpy as np
 
-from jax._src import deprecations
 from jax._src import dtypes
 from jax._src import config
 from jax._src import effects
@@ -638,24 +637,13 @@ def _aval_property(name):
 class Tracer(typing.Array, metaclass=StrictABCMeta):
   __array_priority__ = 1000
   __slots__ = ['_trace', '_line_info']
+  __hash__ = None  # type: ignore
 
   dtype = _aval_property('dtype')
   ndim = _aval_property('ndim')
   size = _aval_property('size')
   shape = _aval_property('shape')
 
-  def __hash__(self):
-    # TODO(jakevdp) finalize this deprecation and set __hash__ = None
-    # Warning added 2024-06-13
-    if deprecations.is_accelerated('tracer-hash'):
-      raise TypeError(f"unhashable type: {type(self)}")
-    # Use FutureWarning rather than DeprecationWarning because hash is likely
-    # not called directly by the user, so we want to warn at all stacklevels.
-    warnings.warn(
-      f"unhashable type: {type(self)}. Attempting to hash a tracer will lead to an"
-      " error in a future JAX release.", category=FutureWarning)
-    return super().__hash__()
-
   def __init__(self, trace: Trace):
     self._trace = trace
 


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index e7aad59b1..9618a8cf4 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -24,7 +24,6 @@ import numpy as np
 import jax
 import jax.numpy as jnp
 from jax._src import core
-from jax._src import deprecations
 from jax._src import dispatch
 from jax._src import op_shardings
 from jax._src import test_util as jtu
@@ -608,16 +607,11 @@ class JaxArrayTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(TypeError, "unhashable type"):
       hash(x)
 
-    @jax.jit
-    def check_tracer_hash(x):
-      self.assertIsInstance(hash(x), int)
+    with self.assertRaisesRegex(TypeError, "unhashable type"):
+      jax.jit(hash)(x)
 
-    if deprecations.is_accelerated('tracer-hash'):
-      with self.assertRaisesRegex(TypeError, "unhashable type"):
-        check_tracer_hash(x)
-    else:
-      with self.assertWarnsRegex(FutureWarning, "unhashable type"):
-        check_tracer_hash(x)
+    with self.assertRaisesRegex(TypeError, "unhashable type"):
+      jax.vmap(hash)(x)
 
   def test_shape_dtype_struct_sharding_jit(self):
     mesh = jtu.create_mesh((8,), ('x'))