Arg<TF_Int32Tensor, [{Groups between which the reductions are performed.}]>:$group_assignment,

    TF_AnyStrAttrOf<["Min", "Max", "Mul", "Add", "Mean"]>:$reduce_op,
    TF_AnyStrAttrOf<["CrossReplica", "CrossReplicaAndPartition"]>:$mode
  );

  let results = (outs
    TensorOf<[TF_Bfloat16, TF_Float16, TF_Float32, TF_Int32, TF_Uint32]>:$output
  );