Remove scheduling annotations on the cloned computation for collective groups

[chaserileyroberts]

Previously, when users try to combine using the collective groups with the scheduling groups, you could run into issues like this:

Example code

@jax.jit
def bidir_comms(a):
    b = jax.lax.ppermute(a, "i", perm_up)
    c = jax.lax.ppermute(a, "i", perm_down)
    return b, c

@jax.jit
@partial(shard_map, mesh=mesh, in_specs=P(None, 'i'), out_specs=P(None, 'i'))
def groups(a):
   # Running the collective groups under a scheduling group.
   with set_xla_metadata( _scheduling_group_id='1'):
      with set_xla_metadata(_collectives_group="", inlineable="false"):
          b, c = bidir_comms(a)
    return b + c

Would crash with an error like

jaxlib.xla_extension.XlaRuntimeError: INTERNAL: There is a scheduling group which exceeds the overlap limits. Annotation id: 1. It needs 2 kGpuAsyncStreamCollectives resources, but the limit is 1. 

This is because the instructions within the async computation also included the scheduling annotations. When the LHS would look within this computation for scheduling, it would see all of the communication operations being labeled with the same group and crash from the overlap limit check. Removing this annotation when creating the async instructions fixes this issue.

This is a simple solution for now, but the real crux of the issue is that with set_xla_metadata is applying frontend attributes to every operation within the context, even though we would prefer it to only be on the inner Call operation. We should consider adding a different JAX API that applies attributes to only the call operation created from a jax.jit.