code: def _group_shard(
      ctx,
      x: ir.Value,
      y: ir.Value,
      x_aval_out: core.AbstractValue,
      y_aval_out: core.AbstractValue,
  ) -> tuple[ir.Value, ir.Value]:
      shard_group_id = next(_next_shard_group_id)
  
      unknown_op_sharding = xc.OpSharding()
      unknown_op_sharding.type = xc.OpSharding.Type.UNKNOWN
      unknown_op_sharding.is_shard_group = True
      unknown_op_sharding.shard_group_id = shard_group_id
      unknown_op_sharding.shard_group_type = (
          xc.OpSharding.ShardGroupType.AS
      )
  
      x = mlir.wrap_with_sharding_op(
          ctx, x, x_aval_out, unknown_op_sharding, has_side_effect=True
      )
      y = mlir.wrap_with_sharding_op(
          ctx, y, y_aval_out, unknown_op_sharding, has_side_effect=True
      )
      return x, y
  
file_path: jax/jax/_src/shard_alike.py
imports: from functools import partial
  import itertools
  from jax._src import core
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.dispatch import apply_primitive
  from jax._src.tree_util import tree_flatten, tree_unflatten
  from jax._src.interpreters import batching
  from jax._src.util import safe_zip
  from jax._src.lib import xla_client
  from jax._src.api_util import shaped_abstractify
  from jax._src.lib.mlir import ir
label_desc: Use image node_id 4 for calling a global function with example usage: _group_shard(ctx, x, y, x_aval_out, y_aval_out) and returns: x,  y