Improve exception message for assert_tree_shape_prefix.

PiperOrigin-RevId: 427413341
google-deepmind · Feb 9, 2022 · 37fd407 · 37fd407
1 parent e54cc52
commit 37fd407
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/chex/_src/asserts.py b/chex/_src/asserts.py
@@ -943,9 +943,9 @@ def _assert_prefix_fn(path, leaf):
     if len(shape_prefix) > len(leaf.shape):
       errors.append(
           (f"Tree leaf '{_ai.format_tree_path(path)}' has a shape "
-           f"of length {len(leaf.shape)} (shape={leaf.shape}) which is smaller "
+           f"of length {leaf.ndim} (shape={leaf.shape}) which is smaller "
            f"than the expected prefix of length {len(shape_prefix)} "
-           "(prefix={shape_prefix})."))
+           f"(prefix={shape_prefix})."))
       return
 
     suffix = leaf.shape[:len(shape_prefix)]