XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaIf").Device(DEVICE), XlaCompileOnDemandOp); \
  REGISTER_KERNEL_BUILDER(Name("XlaOptimizationBarrier").Device(DEVICE),       \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaPad")                                       \

      "XlaDot",
      "XlaDotV2",
      "XlaDynamicSlice",
      "XlaDynamicUpdateSlice",
      "XlaEinsum",
      "XlaGather",
      "XlaIf",
      "XlaKeyValueSort",
      "XlaOptimizationBarrier",
      "XlaPad",
      "XlaRecv",
      "XlaReduce",
      "XlaReducePrecision",
      "XlaReduceScatter",
      "XlaReduceWindow",
      "XlaRemoveDynamicDimensionSize",

  func.return %0 : tensor<i1>
}

//===----------------------------------------------------------------------===//
// tf.XlaOptimizationBarrier legalization
//===----------------------------------------------------------------------===//

  TF_DerivedOperandTypeListAttr Targs = TF_DerivedOperandTypeListAttr<0>;
  TF_DerivedResultTypeListAttr Tresults = TF_DerivedResultTypeListAttr<0>;
}

def TF_XlaOptimizationBarrierOp : TF_Op<"XlaOptimizationBarrier", [Pure]> {
  let summary = "Wraps the XLA OptimizationBarrier operator.";

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