[shard-map] relax vmap-of-shmap error with spmd_axis_name

Because Matt W said so

Co-authored-by: Matthew Wiethoff <wiethoff@google.com>

jax/experimental/shard_map.py

@@ -1313,9 +1313,9 @@ def _shard_map_batch(
                    for ax in names} for names, d in zip(in_names, in_dims)]
   spmd_axis_name = trace.spmd_axis_name
   if spmd_axis_name is not None:
-    used = {n for names in in_names for ns in names.values() for n in ns}
-    if set(spmd_axis_name) & used:
-      raise ValueError("vmap spmd_axis_name cannot appear in shard_map in_specs")
+    for names, dim in zip(in_names, in_dims):
+      if set(spmd_axis_name) & set(names.values()):
+        raise ValueError("vmap spmd_axis_name cannot appear in shard_map in_specs")
     new_in_names = [{**ns, d:spmd_axis_name} if d is not batching.not_mapped  # type: ignore
                     else ns for ns, d in zip(new_in_names, in_dims)]
   @as_hashable_function(closure=out_names_thunk)