topology_attr.getValue(), status_or_device_coodinates.value());
  return success();
}

// Move a broadcast into the XLA cluster, converting it to an XlaAllReduce. This
// skips if the element type is not known to be valid for XlaAllReduce.
LogicalResult MoveBroadcastToCluster(OpBuilder& builder,
                                     OpBuilder& inner_builder,

  // CHECK-NEXT:   }) : () -> tensor<f32>
  // CHECK-NEXT:   "tf_device.cluster"() ({
  // CHECK-NEXT:     %[[GROUP:.*]] = "tf.Const"()
  // CHECK-SAME:       [0, 1, 2, 3]
  // CHECK-NEXT:     %[[REDUCED:.*]] = "tf.XlaAllReduce"(%[[ID]], %[[GROUP]]) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x4xi32>) -> tensor<f32>
  // CHECK-NEXT:     "tf.OpA"(%[[REDUCED]]) : (tensor<f32>) -> ()
  tf_device.replicate {n = 4 : i32} {

  let results = (outs
    TF_FpOrComplexTensor:$z
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;

  let hasCanonicalizer = 1;
}

def TF_XlaAllReduceOp : TF_Op<"XlaAllReduce", [Pure, TF_NoConstantFold]> {
  let summary = "Wraps the XLA AllReduce operator";

  let description = [{
documented at https://www.tensorflow.org/xla/operation_semantics#allreduce.
  }];