Repository: google/jax
Issue: jax.tree_util.Partial is not hash-stable in jax>=0.2.22

Following #8101 , `jax.tree_util.Partial.func` is no more hash-stable:

jax==0.2.21
```python
>>> import jax
>>> from jax.tree_util import Partial
>>> from functools import partial
>>> fun = partial(print, 1)
>>> hash(fun)
278539187
>>> hash(Partial(fun, 2).func)
278539187
>>> hash(Partial(fun, 2).func)
278539187
```

latest release
```python
>>> from jax.tree_util import Partial
>>> from functools import partial
>>> fun = partial(print, 1)
>>> hash(fun)
8749850128694
>>> hash(Partial(fun, 2).func)
8750211514691
>>> hash(Partial(fun, 2).func)
8750211514664

# instead, if the function wrapped is just a standard function, everything works fine
>>> hash(Partial(print, 1).func)
56249262
>>> hash(Partial(print, 1).func)
56249262
```

This is annoying because in some code in NetKet, we were wrapping user-passed functions into a `tree_util.Partial` in order to simplify some code, but now this triggers recompilation if the user's supplied function is wrapped in a partial or not.

cc @hawkinsp who committed the changes

Changed files:
jax/_src/tree_util.py, tests/tree_util_test.py

diff --git a/jax/_src/tree_util.py b/jax/_src/tree_util.py
index f5a798a80..3b5295f59 100644
--- a/jax/_src/tree_util.py
+++ b/jax/_src/tree_util.py
@@ -267,6 +267,24 @@ register_pytree_node(
   lambda s, values: collections.defaultdict(s[0], safe_zip(s[1], values)))  # type: ignore[index]
 
 
+
+class _HashableCallableShim:
+  """Object that delegates __call__, __hash__, and __eq__ to another object."""
+  def __init__(self, fun):
+    self.fun = fun
+
+  def __call__(self, *args, **kw):
+    return self.fun(*args, **kw)
+
+  def __hash__(self):
+    return hash(self.fun)
+
+  def __eq__(self, other):
+    if isinstance(other, _HashableCallableShim):
+      return self.fun == other.fun
+    return self.fun == other
+
+
 class Partial(functools.partial):
   """A version of functools.partial that works in pytrees.
 
@@ -318,16 +336,19 @@ class Partial(functools.partial):
   def __new__(klass, func, *args, **kw):
     # In Python 3.10+, if func is itself a functools.partial instance,
     # functools.partial.__new__ would merge the arguments of this Partial
-    # instance with the arguments of the func. We box func in another lambda to
-    # avoid this optimization since it would change which arguments are
-    # considered part of the pytree.
+    # instance with the arguments of the func. We box func in a class that does
+    # not (yet) have a `func` attribute to defeat this optimization, since we
+    # care exactly which arguments are considered part of the pytree.
     if isinstance(func, functools.partial):
       original_func = func
-      func = lambda *args, **kw: original_func(*args, **kw)
+      func = _HashableCallableShim(original_func)
+      out = super(Partial, klass).__new__(klass, func, *args, **kw)
       func.func = original_func.func
       func.args = original_func.args
       func.keywords = original_func.keywords
-    return super(Partial, klass).__new__(klass, func, *args, **kw)
+      return out
+    else:
+      return super(Partial, klass).__new__(klass, func, *args, **kw)
 
 
 register_pytree_node(
diff --git a/tests/tree_util_test.py b/tests/tree_util_test.py
index 436f9461d..b64b4758a 100644
--- a/tests/tree_util_test.py
+++ b/tests/tree_util_test.py
@@ -200,6 +200,16 @@ class TreeTest(jtu.JaxTestCase):
     h = tree_util.Partial(g, 3)
     self.assertEqual(h.args, (3,))
 
+  def testPartialFuncAttributeHasStableHash(self):
+    # https://github.com/google/jax/issues/9429
+    fun = functools.partial(print, 1)
+    p1 = tree_util.Partial(fun, 2)
+    p2 = tree_util.Partial(fun, 2)
+    self.assertEqual(fun, p1.func)
+    self.assertEqual(p1.func, fun)
+    self.assertEqual(p1.func, p2.func)
+    self.assertEqual(hash(p1.func), hash(p2.func))
+
   @parameterized.parameters(*(TREES + LEAVES))
   def testRoundtripViaBuild(self, inputs):
     xs, tree = _process_pytree(tuple, inputs)