}


// -----

// CHECK: module attributes
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_key = 1
// CHECK-SAME{LITERAL}: tf2xla.collective_info.group_size = 2
// CHECK-LABEL: func @collective_reduce_v2
func.func @collective_reduce_v2(%input: tensor<f32>) -> tensor<f32> {
  %group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>

          StringRef(kGroupKeyAttrName.data(), kGroupKeyAttrName.size()),
          group_key, op, module))) {
    return failure();
  }
  return success();
}

LogicalResult SetCollectiveInfo(OpBuilder& builder,
                                DenseIntElementsAttr replica_groups,
                                Operation* op) {
  // Use special group_key 0 to represent "all available devices". This

      << " is required for CollectiveInfo but not found.";
  int32_t group_key = group_key_attr.getInt();
  int32_t group_size = group_size_attr.getInt();
  VLOG(2) << "Populating CollectiveInfo: group_key=" << group_key
          << " group_size=" << group_size;
  compilation_result->collective_info = {group_key, group_size, 0};
  return absl::OkStatus();
}

Status PopulateResultIOInfo(

  Pattern<
    (TF_CollectiveReduceV2Op:$src_op
       $input, $group_size, $group_key, $instance_key, $ordering_token,
       $merge_op, $final_op, $communication_hint, $timeout_seconds,
       $is_stateless, $max_subdivs_per_device
    ),
    [(TF_CollectiveReduceV2Op
       $input, $group_size, $group_key, $instance_key, (EmptyList),
       $merge_op, $final_op, $communication_hint, $timeout_seconds,

  %input: tensor<f32>,
  %group_key: tensor<i32>,
  %group_size: tensor<i32>,
  %instance_key: tensor<i32>) {
  tf_executor.graph {
    // expected-remark@above {{ID: 5}}
    %island = tf_executor.island {
        // expected-remark@above {{ID: 3}}
        // expected-remark@above {{Successors: {4}}}

      /// $group_assignment
      CompileTimeConstantOperand<TF::CollectiveAssignGroupV2Op, 0>,
      // $source_target_pairs
      CompileTimeConstantOperand<TF::CollectivePermuteOp, 1>,
      // $group_size, $group_key
      CompileTimeConstantOperand<TF::CollectiveReduceV2Op, 1, 2>,
      CompileTimeConstantOperand<TF::ConcatV2Op, -1>,  // (variadic) $axis
      // $filter_sizes

  %group_size: tensor<i32>,
  %group_key: tensor<i32>,
  %instance_key: tensor<i32>,
  %token: tensor<!tf_type.resource<tensor<f32>>>
) -> tensor<3xi64> {
  %0 = "tf_device.cluster"() ({
    // arg4 shall have been removed.
    // CHECK: tf.CollectiveReduceV2
    // CHECK-NOT: [[TOKEN]]
    // CHECK-SAME: merge_op
    %0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key, %token) {

  TF_DeleteStatus(status);

  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(in));
  TFE_OpSetAttrInt(op, "group_size", group_size);
  TFE_OpSetAttrInt(op, "group_key", 123);
  TFE_OpSetAttrInt(op, "instance_key", 456);
  TFE_OpSetAttrString(op, "merge_op", "Add", 3);
  TFE_OpSetAttrString(op, "final_op", "Id", 2);
  std::vector<int64_t> subdiv_offsets;

    "summary": "Instance 127.0.0.1:9000 has lot quorum on pool 0 on set 0"
  },
  "externalURL": "http://fedora-minio:9093",
  "version": "4",
  "groupKey": "{}:{alertname=\"MinIOClusterTolerance\"}",
  "truncatedAlerts": 0
}
```

  );

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

def TF_CollectiveBcastRecvOp : TF_Op<"CollectiveBcastRecv", []> {
  let summary = "Receives a tensor value broadcast from another device.";

  let arguments = (ins
    I64Attr:$group_size,
    I64Attr:$group_key,
    I64Attr:$instance_key,
    TF_ShapeAttr:$shape,