Repository: google/jax

Issue:
dict out_axes in xmap give a hash error
```
def f(x):
  return {"a": x}

x = jax.numpy.ones((2, 3))
xmap(f, in_axes=[...], out_axes={"a": [...]})(x)
```

gives

```
[...]
jax/linear_util.py in memoized_fun(fun, *args)
    253     else:
    254       key = (fun.transforms, fun.params, args, config.x64_enabled)
--> 255     result = cache.get(key, None)
    256     if result is not None:
    257       ans, stores = result

jax/_src/util.py in __hash__(self)
    375 
    376   def __hash__(self):
--> 377     return hash((self.f.__code__, self.closure))
    378 
    379   def __call__(self, *args, **kwargs):

jax/_src/util.py in __hash__(self)
    375 
    376   def __hash__(self):
--> 377     return hash((self.f.__code__, self.closure))
    378 
    379   def __call__(self, *args, **kwargs):

TypeError: unhashable type: 'dict'
```

(The closure of out_axes_thunk seems to contain the out_axes as an unhashable dict!)

Making out_axes a dict-less prefix works:

```
def f(x):
  return {"a": x}

x = jax.numpy.ones((2, 3))
xmap(f, in_axes=[...], out_axes=[...])(x)
> no error
```




Discussion:

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1627,9 +1627,14 @@ def f_pmapped(*args, **kwargs):
         lambda: (0,) * out_tree().num_leaves,
         closure=out_axes)
     else:
+      # out_axes_thunk closes over the out_axes, they are flattened here to make
+      # them hashable.
+      out_axes_leaves, out_axes_treedef = tree_flatten(out_axes)
       out_axes_thunk = HashableFunction(
-        lambda: tuple(flatten_axes("pmap out_axes", out_tree(), out_axes)),
-        closure=out_axes)
+        lambda: tuple(flatten_axes("pmap out_axes", out_tree(),
+                                   tree_unflatten(out_axes_treedef,
+                                                  list(out_axes_leaves)))),
+        closure=(tuple(out_axes_leaves), out_axes_treedef))
     out = pxla.xla_pmap(
         flat_fun, *args, backend=backend, axis_name=axis_name,
         axis_size=local_axis_size, global_axis_size=axis_size,
diff --git a/jax/experimental/maps.py b/jax/experimental/maps.py
--- a/jax/experimental/maps.py
+++ b/jax/experimental/maps.py
@@ -488,9 +488,15 @@ def fun_mapped(*args):
     # TODO: Check that:
     #         - two axes mapped to the same resource never coincide (even inside f)
     in_axes_flat = flatten_axes("xmap in_axes", in_tree, in_axes)
+
+    # out_axes_thunk closes over the out_axes, they are flattened here to make
+    # them hashable.
+    out_axes_leaves, out_axes_treedef = tree_flatten(out_axes)
     out_axes_thunk = HashableFunction(
-      lambda: tuple(flatten_axes("xmap out_axes", out_tree(), out_axes)),
-      closure=out_axes)
+      lambda: tuple(flatten_axes("xmap out_axes", out_tree(),
+                                 tree_unflatten(out_axes_treedef,
+                                                list(out_axes_leaves)))),
+      closure=(tuple(out_axes_leaves), out_axes_treedef))
 
     axis_resource_count = _get_axis_resource_count(normalized_axis_resources, resource_env)
     for axis, size in axis_sizes.items():


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -2033,6 +2033,15 @@ def f(x, y):
     self.assertAllClose(f(x, y),
                         (jnp.sin(x.transpose((1, 0, 2)) + y).transpose((1, 2, 0)), y * 2))
 
+  def testPmapDictOutAxes(self):
+    # see issue #6410
+    @partial(pmap, out_axes={'a': 0})
+    def f(x):
+      return {'a': x}
+    device_count = xla_bridge.device_count()
+    x = jnp.arange(device_count)
+    tree_util.tree_multimap(self.assertAllClose, f(x), {'a': x})
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": f"_{in_axes}_{out_axes}",
        "in_axes": in_axes, "out_axes": out_axes}
diff --git a/tests/xmap_test.py b/tests/xmap_test.py
--- a/tests/xmap_test.py
+++ b/tests/xmap_test.py
@@ -1142,6 +1142,12 @@ def testNegativeAxes(self):
     with self.assertRaisesRegex(ValueError, "xmap doesn't support negative axes in out_axes"):
       xmap(lambda x: x, in_axes={0: 'i'}, out_axes={-1: 'i'})(jnp.ones((5,)))
 
+  @ignore_xmap_warning()
+  def testDictOutAxes(self):
+    # see issue #6410
+    out = xmap(lambda x: x, in_axes=[...], out_axes={"a": [...]})({"a": 1})
+    self.assertEqual(out, {"a": 1})
+
   @ignore_xmap_warning()
   def testListAxesRankAssertion(self):
     error = (r"xmap argument has an in_axes specification of \['i', None\], which "