REGISTER_KERNEL_BUILDER(Name("XlaRecv").Device(DEVICE),                      \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaReduce").Device(DEVICE),                    \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaVariadicReduce").Device(DEVICE),            \

                                         "XlaKeyValueSort",
                                         "XlaPad",
                                         "XlaRecv",
                                         "XlaReduce",
                                         "XlaReduceWindow",
                                         "XlaReplicaId",
                                         "XlaRngBitGenerator",

      "XlaDynamicUpdateSlice",
      "XlaEinsum",
      "XlaGather",
      "XlaIf",
      "XlaKeyValueSort",
      "XlaOptimizationBarrier",
      "XlaPad",
      "XlaRecv",
      "XlaReduce",
      "XlaReducePrecision",
      "XlaReduceScatter",
      "XlaReduceWindow",
      "XlaRemoveDynamicDimensionSize",
      "XlaReplicaId",
      "XlaRngBitGenerator",
      "XlaScatter",

  %0 = "tf.XlaReduce"(%arg0, %arg1) {dimensions_to_reduce = [], reducer = @sum1, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<*xbf16>, tensor<*xbf16>) -> tensor<*xbf16>
  func.return %0 : tensor<*xbf16>
}

  }];

  let arguments = (ins
    StrAttr:$config
  );

  let results = (outs
    TF_VariantTensor:$output
  );
}

def TF_XlaReduceOp : TF_Op<"XlaReduce", [Pure]> {
  let summary = "Wraps the XLA Reduce operator, documented at";

  let description = [{
https://www.tensorflow.org/performance/xla/operation_semantics#reduce .
  }];