Arg<TF_Int32Tensor, [{An int32 tensor with shape
[num_groups, num_replicas_per_group]. `group_assignment[i]` represents the
replica ids in the ith subgroup.}]>:$group_assignment,

    I64Attr:$concat_dimension,
    I64Attr:$split_dimension,
    I64Attr:$split_count
  );

  let results = (outs