let summary = "Assign group keys based on group assignment.";

  let arguments = (ins
    TF_Int32Tensor:$group_assignment,
    TF_Int32Tensor:$device_index,
    TF_Int32Tensor:$base_key
  );

  let results = (outs
    TF_Int32Tensor:$group_size,
    TF_Int32Tensor:$group_key
  );
}

def TF_CollectiveBcastRecvOp : TF_Op<"CollectiveBcastRecv", []> {