std::optional<absl::string_view> tf_program = std::nullopt,
      double error_tolerance = 0.1) {
    // Expected result is calculated by evaluating using TF kernels. In some
    // cases, TF kernel behaves differently from lowered graph (e.g. Hybrid
    // ops). So we optionally use a different graph to calculate the expected
    // result.
    TF_ASSERT_OK_AND_ASSIGN(
        auto expected,

    .Attr("T: type")
    .Attr("enabled: bool")
    .Attr("func_name: string")
    .Attr("node_name: string")
    .SetIsStateful();

class DumpTensorOp : public OpKernel {
 public:
  explicit DumpTensorOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    string log_dir_path;
    string file_name;
    string func_name;
    string node_name;
    OP_REQUIRES_OK(ctx, ctx->GetAttr("log_dir_path", &log_dir_path));

  }

  for (auto hashtable : hashtables) {
    auto key_dtype = hashtable.getKeyDtype();
    auto value_dtype = hashtable.getValueDtype();

    // Only allow string -> int64 and int64 -> string mappings due to kernel
    // capability.
    if (!((mlir::isa<TF::StringType>(key_dtype) &&
           mlir::isa<IntegerType>(value_dtype) &&
           mlir::cast<IntegerType>(value_dtype).getWidth() == 64) ||

    return status;
  }
  devices->push_back(std::move(device));
  return absl::OkStatus();
}

REGISTER_LOCAL_DEVICE_FACTORY(DEVICE_XLA_CPU, XlaCpuDeviceFactory);

// Kernel registrations

constexpr std::array<DataType, 18> kAllXlaCpuTypes = {
    {DT_UINT8, DT_QUINT8, DT_UINT16, DT_INT8, DT_QINT8, DT_INT16, DT_INT32,
     DT_QINT32, DT_INT64, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64,

      c->set_output(3, c->MakeShape({num_bins_attr->i()}));

      return absl::OkStatus();
    });

class CustomAggregatorOp : public OpKernel {
 public:
  explicit CustomAggregatorOp(OpKernelConstruction* context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("id", &id_));

    int calibration_method_value;
    int num_bins;
    float min_percentile;

     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
  let summary = "Copy the CPU fallback tensor if it is small";

  let description = [{
    This kernel performs deep copy on the input tensor if it is small, to avoid
    atomic contention on its refcount.

    Note that this kernel always create a new AsyncValue for each result to
    avoid atomic contention on AsyncValue's refcount.
  }];

  let arguments = (ins

  mlir::LogicalResult PrepareParams();

  // Given the required_consts, it will fill the 3 output vectors with
  // their respective data.
  // Expressions: Output XLA expressions as required by the compiled kernel.
  // Tensors: Vector of tensors that back the TensorValue inputs
  // Inputs: Vector of inputs that are backed by tensors.
  mlir::LogicalResult PrepareKernelInputs(
      const llvm::SmallDenseSet<int>& required_consts,

// e.g. `Add`, that expects its inputs in device memory. Here is how it
// works now.
// First, what do we mean by "op expects an input in XYZ memory"?
// There are two types of "ops" here: the tf2xla kernel and the HLO
// computation it builds. The tf2xla kernel needs to retrieve the actual
// numeric value of the compile-time constant tensors, so it really expects
// them to be on in host memory. However, for other inputs, it refers to them

calibration method and save the result to the given file path as a binary
proto file.)doc");

class CalibrationStatisticsSaverOp : public OpKernel {
 public:
  explicit CalibrationStatisticsSaverOp(
      absl::Nonnull<OpKernelConstruction*> context)
      : OpKernel(context) {
    std::string output_file_path;
    OP_REQUIRES_OK(context,
                   context->GetAttr("output_file_path", &output_file_path));

}

def BatchFunctionOp : TensorflowMlrt_Op<"batch_function", [Pure]> {
  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";

  let description = [{
    This kernel executes a variant tf.BatchFunction kernel that supports having
    the `f` attribute as a bytecode function.

    Example:
      %res = tf_mlrt.batch_function(%input, %captured_input)  {
          device = "/device:CPU:0",