tree_util error handling throws error

[Sohl-Dickstein]

Description

I believe this and this line, should instead be:

        f"    {' '.join(str(k) for k in prefix_tree_keys)}\n"
...
        f"    {' '.join(str(k) for k in full_tree_keys)}\n"

I believe this, because I recently experienced the following error thrown inside this error handling:

...
  File "/root/.pyenv/versions/3.11.6/lib/python3.11/site-packages/jax/_src/pjit.py", line 1346, in _check_and_canonicalize_out_shardings
    out_shardings_flat = flatten_axis_resources(
                         ^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.pyenv/versions/3.11.6/lib/python3.11/site-packages/jax/_src/pjit.py", line 1121, in flatten_axis_resources
    raise e(what)
          ^^^^^^^
  File "/root/.pyenv/versions/3.11.6/lib/python3.11/site-packages/jax/_src/tree_util.py", line 1260, in <lambda>
    f"    {\' \'.join(str(k.key) for k in prefix_tree_keys)}\
"
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.pyenv/versions/3.11.6/lib/python3.11/site-packages/jax/_src/tree_util.py", line 1260, in <genexpr>
    f"    {\' \'.join(str(k.key) for k in prefix_tree_keys)}\
"
                        ^^^^^
AttributeError: \'GetAttrKey\' object has no attribute \'key\'

System info (python version, jaxlib version, accelerator, etc.)

not relevant

[mattjj]

@Sohl-Dickstein 😄 thanks for the fix! I’ll try to apply it tonight when I can.