Raises an explicit error in reshard.

Hit this while working with sharding in types -- passing a sharding that had an empty mesh. (I think this was in a test). This failed trying to acces with `with_spec` attribute on None -- so just catching this case early.

PiperOrigin-RevId: 762135310

Author: jkr26

jax/_src/pjit.py

@@ -2987,6 +2987,11 @@ def reshard(xs, out_shardings):
   out_flat = []
   for x, x_aval, s in safe_zip(x_flat, x_avals_flat, shardings_flat):
     ds = canonicalize_sharding(s, 'reshard')
+    if ds is None:
+      raise ValueError(
+          'Reshard should only be used with out_shardings which are non-None '
+          'and have a nonempty mesh. Got sharding {s}.'
+      )
     ds = ds.with_spec(ds.spec._normalized_spec_for_aval(x_aval.ndim))  # pytype: disable=attribute-error
     out_flat.append(reshard_p.bind(x, dst_sharding=ds))
   return tree_unflatten(treedef, out_flat)

tests/pjit_test.py

@@ -8781,6 +8781,16 @@ def f(x):
                                 "to Shardy"):
       jax.jit(f)(x)
 
+  def test_reshard_empty_mesh_error(self):
+    arr = jax.device_put(np.arange(8), jax.devices()[0])
+    with self.assertRaisesRegex(ValueError, "nonempty mesh"):
+      reshard(arr, NamedSharding(mesh_lib.empty_abstract_mesh, P(None)))
+
+  def test_reshard_none_sharding_error(self):
+    arr = jax.device_put(np.arange(8), jax.devices()[0])
+    with self.assertRaisesRegex(ValueError, "non-None"):
+      reshard(arr, None)
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())