Bad error message when shardings are non-divisible

[yashk2810]

Description

Repro:

In [11]: mesh = jax.make_mesh((4, 2), ('x', 'y'))

In [12]: @partial(jax.jit, out_shardings=s)
    ...: def f(x):
    ...:     return x
    ...:

In [13]: s = NamedSharding(mesh, P('x'))

In [14]: f(jnp.arange(5))
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[14], line 1
----> 1 f(jnp.arange(5))

    [... skipping hidden 9 frame]

File ~/venv/lib/python3.12/site-packages/jax/_src/pjit.py:1417, in pjit_check_aval_sharding(shardings, flat_avals, names, what_aval, allow_uneven_sharding)
   1415 for i, size in enumerate(num_ways_dim_sharded):
   1416   if not allow_uneven_sharding and shape[i] % size != 0:
-> 1417     raise ValueError(f"One of {what_aval}{name_str} was given the sharding "
   1418                      f"of {s}, which implies that "
   1419                      f"the global size of its dimension {i} should be "
   1420                      f"divisible by {size}, but it is equal to {shape[i]} "
   1421                      f"(full shape: {shape})")

ValueError: One of pjit outputs with pytree key path result was given the sharding of NamedSharding(mesh=Mesh('x': 4, 'y': 2, axis_types=(Auto, Auto)), spec=PartitionSpec('x',), memory_kind=unpinned_host), which implies that the global size of its dimension 0 should be divisible by 4, but it is equal to 5 (full shape: (5,))

This sentence in the error message: pytree key path result is confusing because there is no name given for the output. So we should just not print this as it leads to more confusion.

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

N/A