Repository: google/jax
Issue: Computation with distributed arrays results in unexpected PartitionSpec

### Description

Computation with distributed arrays seems to normalize `PartitionSpec` in at least three ways:

- Leading axes of length 1 are replaced by `None`.
- Trailing axes of length 1 are removed.
- Strings are replaced by tuples of strings.

I found this behavior rather unexpected, because this normalization is applied and validated somewhat inconsistently:
1. `PartitionSpec` objects don't get normalized when put into a `NamedSharding`.
1. Partition specs on arrays created with `device_put` also are not get normalized.
2. Differently normalized `NamedSharding` objects sometimes (but not always) compare equal.
3. Differently normalized `PartitionSpec` objects never compare equal.

Example code:
```python
import chex
import jax
import jax.numpy as jnp
import numpy as np

chex.set_n_cpu_devices(8)
P = jax.sharding.PartitionSpec

mesh = jax.sharding.Mesh(jax.devices(), axis_names=['x'])
spec = P('x')
sharding = jax.sharding.NamedSharding(mesh, spec)

x = jax.device_put(jnp.arange(32), sharding)
y = x + 1

print('example one:')
print(f'{x.sharding=}')
print(f'{y.sharding=}')
print(x.sharding == y.sharding)
print(x.sharding.spec == y.sharding.spec)

devices = np.reshape(jax.devices(), (1, 8, 1))
mesh = jax.sharding.Mesh(devices, axis_names=['x', 'y', 'z'])
spec = P('x', 'y', 'z')
sharding = jax.sharding.NamedSharding(mesh, spec)

x = jax.device_put(jnp.arange(32).reshape(1, -1, 1), sharding)
y = x + 1

print()
print('example two:')
print(f'{x.sharding=}')
print(f'{y.sharding=}')
print(x.sharding == y.sharding)
print(x.sharding.spec == y.sharding.spec)
```
Prints:
```
example one:
x.sharding=NamedSharding(mesh={'x': 8}, spec=PartitionSpec('x',))
y.sharding=NamedSharding(mesh={'x': 8}, spec=PartitionSpec(('x',),))
True
False

example two:
x.sharding=NamedSharding(mesh={'x': 1, 'y': 8, 'z': 1}, spec=PartitionSpec('x', 'y', 'z'))
y.sharding=NamedSharding(mesh={'x': 1, 'y': 8, 'z': 1}, spec=PartitionSpec(None, ('y',)))
False
False
```

### What jax/jaxlib version are you using?

_No response_

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

Changed files:
tests/pjit_test.py, jax/_src/sharding_impls.py, jax/_src/interpreters/pxla.py, jax/_src/pjit.py

diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index 2d8cd11fc..4de5f7446 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -2435,13 +2435,18 @@ orig_out_sharding_handlers[sharding_impls.PositionalSharding] = _gspmd_to_positi
 
 
 def _get_out_sharding_from_orig_sharding(
-    out_shardings, orig_s, are_out_sharding_from_xla):
+    out_shardings, orig_s, orig_aval, are_out_sharding_from_xla):
   out = []
   orig_handler = orig_out_sharding_handlers[type(orig_s)]
   for o, from_xla in safe_zip(out_shardings, are_out_sharding_from_xla):
     if isinstance(o, sharding_impls.GSPMDSharding):
       try:
-        out.append((orig_handler(o._op_sharding, orig_s), False))
+        if (orig_aval is not None and
+            sharding_impls.are_op_shardings_equal(
+                o._op_sharding, orig_s._to_xla_op_sharding(orig_aval.ndim))):
+          out.append((orig_s, False))
+        else:
+          out.append((orig_handler(o._op_sharding, orig_s), False))
       except:
         out.append((o, from_xla))
     else:
@@ -2449,20 +2454,22 @@ def _get_out_sharding_from_orig_sharding(
   return out
 
 def maybe_get_orig_out_sharding(
-    in_shardings, out_shardings, are_out_shardings_from_xla):
+    in_shardings, out_shardings, are_out_shardings_from_xla, in_avals):
   if all(hasattr(o, '_original_sharding') for o in out_shardings):
     return ([o._original_sharding for o in out_shardings],
             (False,) * len(out_shardings))
 
   orig_s = None
-  for i in in_shardings:
+  orig_aval = None
+  for i, aval in safe_zip(in_shardings, in_avals):
     oi = getattr(i, '_original_sharding', None)
     if type(oi) in orig_out_sharding_handlers:
       orig_s = oi
+      orig_aval = aval
       break
   if orig_s is not None:
     return zip(*_get_out_sharding_from_orig_sharding(
-        out_shardings, orig_s, are_out_shardings_from_xla))
+        out_shardings, orig_s, orig_aval, are_out_shardings_from_xla))
 
   return out_shardings, are_out_shardings_from_xla
 
@@ -2675,7 +2682,8 @@ class UnloadedMeshExecutable:
           xla_executable.local_devices(), len(in_shardings), len(out_shardings))
 
     out_shardings, are_out_shardings_from_xla = maybe_get_orig_out_sharding(
-        in_shardings, out_shardings, are_out_shardings_from_xla)
+        in_shardings, out_shardings, are_out_shardings_from_xla,
+        global_in_avals)
 
     return UnloadedMeshExecutable(
         xla_executable=xla_executable,
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index ba281861c..ea23e2e4d 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -1374,7 +1374,7 @@ def _pjit_batcher_for_sharding(
     new_gs = GSPMDSharding(s._device_assignment, new_op)  # type: ignore
     if hasattr(s, '_original_sharding'):
       vmapped_s, _ = pxla._get_out_sharding_from_orig_sharding(
-          [new_gs], s._original_sharding, [False])[0]  # type: ignore
+          [new_gs], s._original_sharding, None, [False])[0]  # type: ignore
       new_gs = to_gspmd_sharding(vmapped_s, ndim)
     return new_gs
   else:
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index 003daf56e..46386b6a8 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -1167,6 +1167,7 @@ def parse_flatten_op_sharding(op_sharding: xc.OpSharding,
         raise NotImplementedError("Unhandled OpSharding type. Please open a bug report!")
     if replicate_on_last_tile_dim:
       partitions = partitions[:-1]
-    return [ParsedPartitionSpec('<internally generated spec>', partitions)]
+    return [CanonicalizedParsedPartitionSpec(
+        ParsedPartitionSpec('<internally generated spec>', partitions))]
   else:
     raise AssertionError("Unhandled OpSharding type. Please open a bug report!")
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index da455227b..9016610fc 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -3311,6 +3311,13 @@ class ArrayPjitTest(jtu.JaxTestCase):
     # and 2 in _array_shard_arg.
     self.assertEqual(cache_info2.misses, cache_info1.misses + 4)
 
+  def test_same_named_sharding_pspec_on_eager_ops(self):
+    mesh = jtu.create_global_mesh((1, 8, 1), ('x', 'y', 'z'))
+    sharding = jax.sharding.NamedSharding(mesh, P('x', 'y', 'z'))
+    x = jax.device_put(jnp.arange(32).reshape(1, -1, 1), sharding)
+    y = x + 1
+    self.assertEqual(x.sharding, y.sharding)
+
 
 class TempSharding(Sharding):
 
@@ -3660,6 +3667,8 @@ class UtilTest(jtu.JaxTestCase):
       spec = [()] * dims
       for axis in rng.permutation(mesh_axes)[:rng.randint(low=1, high=len(mesh_axes) + 1)]:
         spec[rng.choice(dims)] += (axis,)
+      while spec and spec[-1] == ():
+        spec.pop()
       roundtrip(P(*spec))
 
   @parameterized.named_parameters(