REGISTER_KERNEL_BUILDER(Name("XlaSharding").Device(DEVICE),                  \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaReplicaId").Device(DEVICE),                 \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(                                                     \

                 )
                ),
                (replaceWithValue $output)
              ]>;

//===----------------------------------------------------------------------===//
// XlaReplicaId op.
//===----------------------------------------------------------------------===//

def : Pat<(TF_XlaReplicaIdOp),
          (TF_CastOp (MHLO_ReplicaIdOp), /*truncate=*/ConstBoolAttrFalse)>;

                                         "XlaRecv",
                                         "XlaReduce",
                                         "XlaReduceWindow",
                                         "XlaReplicaId",
                                         "XlaRngBitGenerator",
                                         "XlaScatter",
                                         "XlaSelectAndScatter",

      "XlaPad",
      "XlaRecv",
      "XlaReduce",
      "XlaReducePrecision",
      "XlaReduceScatter",
      "XlaReduceWindow",
      "XlaRemoveDynamicDimensionSize",
      "XlaReplicaId",
      "XlaRngBitGenerator",
      "XlaScatter",
      "XlaSelectAndScatter",
      "XlaSelfAdjointEig",
      "XlaSend",
      "XlaSetBound",
      "XlaSetDynamicDimensionSize",

}

func.func @replica_id() -> tensor<i32> {
  // CHECK: %[[ID:.*]] = mhlo.replica_id : tensor<ui32>
  // CHECK: %[[RESULT:.*]] = mhlo.convert %0 : (tensor<ui32>) -> tensor<i32>
  %0 = "tf.XlaReplicaId"() : () -> tensor<i32>
  func.return %0 : tensor<i32>
}

// CHECK: func @angle_c64
// CHECK-SAME: ([[ARG0:%.*]]: tensor<complex<f32>>)
func.func @angle_c64(%arg0: tensor<complex<f32>>) -> tensor<f32> {

  );

  let results = (outs
    TF_Tensor:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_XlaReplicaIdOp : TF_Op<"XlaReplicaId", [Pure, TF_NoConstantFold]> {
  let summary = "Replica ID.";

  let arguments = (ins);

  let results = (outs
    TF_Int32Tensor:$id
  );