Repository: google/jax

Issue:
Better `pjit in_axis_resources specification must be a tree prefix of the corresponding value` error
When using pjit with large, complicated pytrees of arguments and partition specification, the runtime error `ValueError: pjit in_axis_resources specification must be a tree prefix of the corresponding value, got specification` is hard to work with, because it outputs the pspec pytree, and the value pytree (with asterisks) as huge single-line blobs.

An easy first step towards making these errors easier to work with would be to print out the pytrees like JSONs, with linebreaks and multiple indentation levels.

An more involved improvement would be to point out the difference between the two trees to the user in an easy-to-work-with way. For example, the smallest differing sub-pytrees could be printed, together with their locations in the main pytree (as list of keys/indices `(k_i)` that one needs to select in order, to traverse from the root down to the differing sub-pytree: `sub_pytree == root[k_1][...][k_n]`).


Discussion:

Resolving patch:
diff --git a/jax/_src/tree_util.py b/jax/_src/tree_util.py
--- a/jax/_src/tree_util.py
+++ b/jax/_src/tree_util.py
@@ -13,11 +13,13 @@
 # limitations under the License.
 
 import collections
+import difflib
 import functools
 from functools import partial
 import operator as op
 from typing import (Any, Callable, Hashable, Iterable, Optional, Tuple, List,
                     Dict, Type, TypeVar, overload, TYPE_CHECKING, NamedTuple)
+import textwrap
 
 from jax._src.lib import pytree
 
@@ -380,9 +382,10 @@ def flatten_one_level(pytree: Any) -> Tuple[List[Any], Hashable]:
   else:
     raise ValueError(f"can't tree-flatten type: {type(pytree)}")
 
-def prefix_errors(prefix_tree: Any, full_tree: Any
+def prefix_errors(prefix_tree: Any, full_tree: Any,
+                  is_leaf: Optional[Callable[[Any], bool]] = None,
                   ) -> List[Callable[[str], ValueError]]:
-  return list(_prefix_error(KeyPath(()), prefix_tree, full_tree))
+  return list(_prefix_error(KeyPath(()), prefix_tree, full_tree, is_leaf))
 
 class KeyPathEntry(NamedTuple):
   key: Any
@@ -437,38 +440,59 @@ def register_keypaths(ty: Type, handler: Callable[[Any], List[KeyPathEntry]]
 register_keypaths(dict,
                   lambda dct: [GetitemKeyPathEntry(k) for k in sorted(dct)])
 
-def _prefix_error(key_path: KeyPath, prefix_tree: Any, full_tree: Any
+def _prefix_error(key_path: KeyPath, prefix_tree: Any, full_tree: Any,
+                  is_leaf: Optional[Callable[[Any], bool]] = None,
                   ) -> Iterable[Callable[[str], ValueError]]:
   # A leaf is a valid prefix of any tree:
-  if treedef_is_leaf(tree_structure(prefix_tree)): return
+  if treedef_is_leaf(tree_structure(prefix_tree, is_leaf=is_leaf)): return
 
   # The subtrees may disagree because their roots are of different types:
   if type(prefix_tree) != type(full_tree):
     yield lambda name: ValueError(
-      "pytree structure error: different types "
-      f"at {{name}}{key_path.pprint()}: "
-      f"prefix pytree {{name}} has type {type(prefix_tree)} "
-      f"where full pytree has type {type(full_tree)}.".format(name=name))
+      "pytree structure error: different types at key path\n"
+      f"    {{name}}{key_path.pprint()}\n"
+      f"At that key path, the prefix pytree {{name}} has a subtree of type\n"
+      f"    {type(prefix_tree)}\n"
+      f"but at the same key path the full pytree has a subtree of different type\n"
+      f"    {type(full_tree)}.".format(name=name))
     return  # don't look for more errors in this subtree
 
-  # Or they may disagree if their roots have different numbers of children:
+  # Or they may disagree if their roots have different numbers of children (note
+  # that because both prefix_tree and full_tree have the same type at this
+  # point, and because prefix_tree is not a leaf, each can be flattened once):
   prefix_tree_children, prefix_tree_meta = flatten_one_level(prefix_tree)
   full_tree_children, full_tree_meta = flatten_one_level(full_tree)
   if len(prefix_tree_children) != len(full_tree_children):
     yield lambda name: ValueError(
-      "pytree structure error: different numbers of pytree children "
-      f"at {{name}}{key_path.pprint()}: "
-      f"prefix pytree {{name}} has {len(prefix_tree_children)} children where "
-      f"full pytree has {len(full_tree_children)} children.".format(name=name))
+      "pytree structure error: different numbers of pytree children at key path\n"
+      f"    {{name}}{key_path.pprint()}\n"
+      f"At that key path, the prefix pytree {{name}} has a subtree of type\n"
+      f"    {type(prefix_tree)}\n"
+      f"with {len(prefix_tree_children)} children, "
+      f"but at the same key path the full pytree has a subtree of the same "
+      f"type but with {len(full_tree_children)} children.".format(name=name))
     return  # don't look for more errors in this subtree
 
   # Or they may disagree if their roots have different pytree metadata:
   if prefix_tree_meta != full_tree_meta:
+    prefix_tree_meta_str = str(prefix_tree_meta)
+    full_tree_meta_str = str(full_tree_meta)
+    metadata_diff = textwrap.indent(
+        '\n'.join(difflib.ndiff(prefix_tree_meta_str.splitlines(),
+                                full_tree_meta_str.splitlines())),
+        prefix="    ")
     yield lambda name: ValueError(
-      "pytree structure error: different pytree metadata "
-      f"at {{name}}{key_path.pprint()}: "
-      f"prefix pytree {{name}} has metadata {prefix_tree_meta} where "
-      f"full pytree has metadata {full_tree_meta}.".format(name=name))
+      "pytree structure error: different pytree metadata at key path\n"
+      f"    {{name}}{key_path.pprint()}\n"
+      f"At that key path, the prefix pytree {{name}} has a subtree of type\n"
+      f"    {type(prefix_tree)}\n"
+      f"with metadata\n"
+      f"    {prefix_tree_meta_str}\n"
+      f"but at the same key path the full pytree has a subtree of the same "
+      f"type but with metadata\n"
+      f"    {full_tree_meta_str}\n"
+      f"so the diff in the metadata at these pytree nodes is\n"
+      f"{metadata_diff}".format(name=name))
     return  # don't look for more errors in this subtree
 
   # If the root types and numbers of children agree, there must be an error
diff --git a/jax/experimental/pjit.py b/jax/experimental/pjit.py
--- a/jax/experimental/pjit.py
+++ b/jax/experimental/pjit.py
@@ -40,7 +40,9 @@
 from jax.interpreters import partial_eval as pe
 from jax.interpreters.sharded_jit import PartitionSpec
 from jax._src.lib import xla_client as xc
-from jax.tree_util import tree_map, tree_flatten, tree_unflatten
+from jax.tree_util import (tree_map, tree_flatten, tree_unflatten,
+                           treedef_is_leaf, tree_structure)
+from jax._src.tree_util import prefix_errors
 from jax._src.util import (extend_name_stack, HashableFunction, safe_zip,
                          wrap_name, wraps, distributed_debug_log,
                          split_list, cache, tuple_insert)
@@ -293,11 +295,60 @@ def flatten_axis_resources(what, tree, axis_resources, tupled_args):
   try:
     return tuple(flatten_axes(what, tree, axis_resources, tupled_args=tupled_args))
   except ValueError:
-    pass
+    pass  # Raise a tree prefix error below
+
+  # Tree leaves are always valid prefixes, so if there was a prefix error as
+  # assumed here, axis_resources must not be a leaf.
+  assert not treedef_is_leaf(tree_structure(axis_resources))
+
+  # Check the type directly rather than using isinstance because of namedtuples.
+  if tupled_args and (type(axis_resources) is not tuple or
+                      len(axis_resources) != len(tree.children())):
+    # We know axis_resources is meant to be a tuple corresponding to the args
+    # tuple, but while it is a non-leaf pytree, either it wasn't a tuple or it
+    # wasn't the right length.
+    msg = (f"{what} specification must be a tree prefix of the positional "
+           f"arguments tuple passed to the `pjit`-decorated function. In "
+           f"particular, {what} must either be a None, a PartitionSpec, or "
+           f"a tuple of length equal to the number of positional arguments.")
+    # If `tree` represents an args tuple, then `axis_resources` must be a tuple.
+    # TODO(mattjj,apaszke): disable implicit list casts, remove 'or list' below
+    if type(axis_resources) is not tuple:
+      msg += f" But {what} is not a tuple: got {type(axis_resources)} instead."
+    elif len(axis_resources) != len(tree.children()):
+      msg += (f" But {what} is the wrong length: got a tuple or list of length "
+              f"{len(axis_resources)} for an args tuple of length "
+              f"{len(tree.children())}.")
+
+    # As an extra hint, let's check if the user just forgot to wrap
+    # in_axis_resources in a singleton tuple.
+    if len(tree.children()) == 1:
+      try: flatten_axes(what, tree, (axis_resources,))
+      except ValueError: pass  # That's not the issue.
+      else:
+        msg += (f" Given the corresponding argument being "
+                f"passed, it looks like {what} might need to be wrapped in "
+                f"a singleton tuple.")
+
+    raise ValueError(msg)
+
   # Replace axis_resources with unparsed versions to avoid revealing internal details
-  flatten_axes(what, tree, tree_map(lambda parsed: parsed.user_spec, axis_resources),
-               tupled_args=tupled_args)
-  raise AssertionError("Please open a bug request!")  # This should be unreachable
+  axis_tree = tree_map(lambda parsed: parsed.user_spec, axis_resources)
+
+  # Because ecause we only have the `tree` treedef and not the full pytree here,
+  # we construct a dummy tree to compare against. Revise this in callers?
+  dummy_tree = tree_unflatten(tree, [PytreeLeaf()] * tree.num_leaves)
+  errors = prefix_errors(axis_tree, dummy_tree)
+  if errors:
+    e = errors[0]  # Only show information about the first disagreement found.
+    raise e(what)
+
+  # At this point we've failed to find a tree prefix error.
+  assert False, "Please open a bug report!"  # This should be unreachable.
+
+class PytreeLeaf:
+  def __repr__(self): return "pytree leaf"
+
 
 @lu.cache
 def _pjit_jaxpr(fun, mesh, local_in_avals,
@@ -472,8 +523,7 @@ def __repr__(self):
 def _prepare_axis_resources(axis_resources,
                             arg_name,
                             allow_unconstrained_dims=False):
-  # PyTrees don't treat None values as leaves, so we explicitly need
-  # to explicitly declare them as such
+  # PyTrees don't treat None values as leaves, so we use an is_leaf function.
   entries, treedef = tree_flatten(axis_resources, is_leaf=lambda x: x is None)
   what = f"{arg_name} leaf specifications"
   entries = [
diff --git a/jax/tree_util.py b/jax/tree_util.py
--- a/jax/tree_util.py
+++ b/jax/tree_util.py
@@ -55,4 +55,7 @@
   treedef_children as treedef_children,
   treedef_is_leaf as treedef_is_leaf,
   treedef_tuple as treedef_tuple,
+  register_keypaths as register_keypaths,
+  AttributeKeyPathEntry as AttributeKeyPathEntry,
+  GetitemKeyPathEntry as GetitemKeyPathEntry,
 )


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -1138,15 +1138,32 @@ def testEmptyMesh(self):
   def testAxisResourcesMismatch(self):
     x = jnp.ones([])
     p = [None, None, None]
+
     pjit(lambda x: x, (p,), p)([x, x, x])  # OK
+
     error = re.escape(
-        r"pjit in_axis_resources specification must be a tree prefix of the "
-        r"corresponding value, got specification (None, None, None) for value "
-        r"tree PyTreeDef((*, *)). Note that pjit in_axis_resources that are "
-        r"non-trivial pytrees should always be wrapped in a tuple representing "
-        r"the argument list.")
+        "pjit in_axis_resources specification must be a tree prefix of the "
+        "positional arguments tuple passed to the `pjit`-decorated function. "
+        "In particular, pjit in_axis_resources must either be a None, a "
+        "PartitionSpec, or a tuple of length equal to the number of positional "
+        "arguments. But pjit in_axis_resources is the wrong length: got a "
+        "tuple or list of length 3 for an args tuple of length 2.")
     with self.assertRaisesRegex(ValueError, error):
-      pjit(lambda x, y: x, p, p)(x, x)  # Error, but make sure we hint at tupling
+      pjit(lambda x, y: x, p, p)(x, x)
+
+    Foo = namedtuple('Foo', ['x'])
+    error = "in_axis_resources is not a tuple.*might need to be wrapped"
+    with self.assertRaisesRegex(ValueError, error):
+      pjit(lambda x: x, Foo(None), Foo(None))(Foo(x))
+
+    pjit(lambda x: x, (Foo(None),), Foo(None))(Foo(x))  # OK w/ singleton tuple
+
+    # TODO(apaszke,mattjj): Disable implicit list casts and enable this
+    # error = ("it looks like pjit in_axis_resources might need to be wrapped in "
+    #          "a singleton tuple.")
+    # with self.assertRaisesRegex(ValueError, error):
+    #   pjit(lambda x, y: x, p, p)([x, x, x])
+
     # TODO(apaszke): Disable implicit list casts and enable this
     # error = re.escape(
     # r"pjit in_axis_resources specification must be a tree prefix of the "
@@ -1158,10 +1175,16 @@ def testAxisResourcesMismatch(self):
     # r"singleton tuple.")
     # with self.assertRaisesRegex(ValueError, error):
     # pjit(lambda x: x, p, p)([x, x, x])  # Error, but make sure we hint at singleton tuple
+
     error = re.escape(
-        r"pjit out_axis_resources specification must be a tree prefix of the "
-        r"corresponding value, got specification [[None, None, None], None] for "
-        r"value tree PyTreeDef([*, *, *]).")
+        "pytree structure error: different numbers of pytree children at "
+        "key path\n"
+        "    pjit out_axis_resources tree root\n"
+        "At that key path, the prefix pytree pjit out_axis_resources has a "
+        "subtree of type\n"
+        "    <class 'list'>\n"
+        "with 2 children, but at the same key path the full pytree has a "
+        "subtree of the same type but with 3 children.")
     with self.assertRaisesRegex(ValueError, error):
       pjit(lambda x: x, (p,), [p, None])([x, x, x])  # Error, we raise a generic tree mismatch message
 
diff --git a/tests/tree_util_test.py b/tests/tree_util_test.py
--- a/tests/tree_util_test.py
+++ b/tests/tree_util_test.py
@@ -439,79 +439,91 @@ class TreePrefixErrorsTest(jtu.JaxTestCase):
 
   def test_different_types(self):
     e, = prefix_errors((1, 2), [1, 2])
-    expected = "pytree structure error: different types at in_axes tree root"
+    expected = ("pytree structure error: different types at key path\n"
+                "    in_axes tree root")
     with self.assertRaisesRegex(ValueError, expected):
       raise e('in_axes')
 
   def test_different_types_nested(self):
     e, = prefix_errors(((1,), (2,)), ([3], (4,)))
-    expected = r"pytree structure error: different types at in_axes\[0\]"
+    expected = ("pytree structure error: different types at key path\n"
+                r"    in_axes\[0\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e('in_axes')
 
   def test_different_types_multiple(self):
     e1, e2 = prefix_errors(((1,), (2,)), ([3], [4]))
-    expected = r"pytree structure error: different types at in_axes\[0\]"
+    expected = ("pytree structure error: different types at key path\n"
+                r"    in_axes\[0\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e1('in_axes')
-    expected = r"pytree structure error: different types at in_axes\[1\]"
+    expected = ("pytree structure error: different types at key path\n"
+                r"    in_axes\[1\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e2('in_axes')
 
   def test_different_num_children(self):
     e, = prefix_errors((1,), (2, 3))
     expected = ("pytree structure error: different numbers of pytree children "
-                "at in_axes tree root")
+                "at key path\n"
+                "    in_axes tree root")
     with self.assertRaisesRegex(ValueError, expected):
       raise e('in_axes')
 
   def test_different_num_children_nested(self):
     e, = prefix_errors([[1]], [[2, 3]])
     expected = ("pytree structure error: different numbers of pytree children "
-                r"at in_axes\[0\]")
+                "at key path\n"
+                r"    in_axes\[0\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e('in_axes')
 
   def test_different_num_children_multiple(self):
     e1, e2 = prefix_errors([[1], [2]], [[3, 4], [5, 6]])
     expected = ("pytree structure error: different numbers of pytree children "
-                r"at in_axes\[0\]")
+                "at key path\n"
+                r"    in_axes\[0\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e1('in_axes')
     expected = ("pytree structure error: different numbers of pytree children "
-                r"at in_axes\[1\]")
+                "at key path\n"
+                r"    in_axes\[1\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e2('in_axes')
 
   def test_different_metadata(self):
     e, = prefix_errors({1: 2}, {3: 4})
     expected = ("pytree structure error: different pytree metadata "
-                "at in_axes tree root")
+                "at key path\n"
+                "    in_axes tree root")
     with self.assertRaisesRegex(ValueError, expected):
       raise e('in_axes')
 
   def test_different_metadata_nested(self):
     e, = prefix_errors([{1: 2}], [{3: 4}])
     expected = ("pytree structure error: different pytree metadata "
-                r"at in_axes\[0\]")
+                "at key path\n"
+                r"    in_axes\[0\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e('in_axes')
 
   def test_different_metadata_multiple(self):
     e1, e2 = prefix_errors([{1: 2}, {3: 4}], [{3: 4}, {5: 6}])
     expected = ("pytree structure error: different pytree metadata "
-                r"at in_axes\[0\]")
+                "at key path\n"
+                r"    in_axes\[0\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e1('in_axes')
     expected = ("pytree structure error: different pytree metadata "
-                r"at in_axes\[1\]")
+                "at key path\n"
+                r"    in_axes\[1\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e2('in_axes')
 
   def test_fallback_keypath(self):
     e, = prefix_errors(Special(1, [2]), Special(3, 4))
-    expected = ("pytree structure error: different types at "
-                r"in_axes\[<flat index 1>\]")
+    expected = ("pytree structure error: different types at key path\n"
+                r"    in_axes\[<flat index 1>\]")
     with self.assertRaisesRegex(ValueError, expected):
       raise e('in_axes')