patterns->add<odml::ConvertCustomCallOp>(context);
  populateWithGenerated(*patterns);

  patterns->add<ConvertReduceOpToTFLiteArgmin, ConvertReduceOpToTFLiteArgmax>(
      context);
}

// Creates an instance of the pass.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeHloToTfLitePass() {
  return std::make_unique<LegalizeHloToTfLitePass>();
}

// Registers the pass implementation

        }
      }

      // Replace the arg with a tf.Const op in the function body.
      builder.setInsertionPointToStart(&func.getBody().front());
      auto const_op = builder.create<TF::ConstOp>(global_tensor.getLoc(),
                                                  *global_tensor.getValue());
      args_to_erase.set(val.getArgNumber());
      for (auto read_op : read_variable_ops_to_erase) {

  return strings::StrCat(forward_frame_name, "_backprop");
}

// Creates a loop that counts the number of iterations performed by the
// while loop associated with `while_ctx`. The returned output yields the
// iteration count.
Status AddForwardLoopCounter(WhileContext* while_ctx, const Scope& scope,
                             Output* count) {
  // Create while loop:
  //   i = 0
  //   while forward loop predicate is true:

      // Add dequantize.
      builder.setInsertionPointAfter(read_variable_op);
      auto new_read_variable_op =
          builder.create<ReadVariableOp>(read_variable_op.getLoc(), ref_qtype,
                                         read_variable_op.getResourceId());
      auto new_dq_op = builder.create<DequantizeOp>(
          read_variable_op.getLoc(), read_variable_op.getResult().getType(),
          new_read_variable_op.getResult());

  // available during compilation or compilation only device for on demand
  // execution which may create a recursion if used for constant folding.
  std::string host_cpu = tensorflow::DeviceNameUtils::FullName(
      /*job=*/"localhost", /*replica=*/0, /*task=*/0, /*type=*/"CPU", /*id=*/0);

  absl::StatusOr<OpKernelRunner> runner = OpKernelRunner::Create(
      node_def->get()->op(), node_def->get()->name(), host_cpu, operands.size(),

    b.setInsertionPoint(op);
    if (auto dq = llvm::dyn_cast<DequantizeOp>(op)) {
      auto dcast = b.create<quantfork::DequantizeCastOp>(
          dq.getLoc(), dq.getOutput().getType(), dq.getInput());
      dq.getOutput().replaceAllUsesWith(dcast);
      dq.erase();
    } else if (auto q = llvm::dyn_cast<QuantizeOp>(op)) {
      auto qcast = b.create<quantfork::QuantizeCastOp>(
          q.getLoc(), q.getOutput().getType(), q.getInput());

  builder.setInsertionPoint(composite);
  if (IsKVCacheCompositeOp(composite)) {
    return builder.create<TFL::CustomOp>(
        composite->getLoc(), composite->getResultTypes(),
        composite->getOperands().slice(2, 3), composite.getName(),
        CustomOption(&builder, custom_option_buffer));
  }
  return builder.create<TFL::CustomOp>(
      composite->getLoc(), composite->getResultTypes(),

namespace mlir {
namespace odml {

LogicalResult ConvertCustomCallOp::matchAndRewrite(
    mhlo::CustomCallOp mhlo_custom_call, OpAdaptor adaptor,
    ConversionPatternRewriter& rewriter) const {
  auto tfl_custom = rewriter.create<TFL::CustomOp>(
      mhlo_custom_call.getLoc(), mhlo_custom_call.getResultTypes(),
      mhlo_custom_call.getInputs());
  tfl_custom.setCustomCodeAttr(
      rewriter.getStringAttr(mhlo_custom_call.getCallTargetName()));

    const llvm::SmallVector<int64_t>& expected_packed_shape,
    const llvm::SmallVector<int8_t>& expected_packed_values) {
  MLIRContext context;
  OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
  OpBuilder builder(&module->getBodyRegion());
  context.loadDialect<TF::TensorFlowDialect>();

  Value value = CreateConstValue<int8_t>(builder, module->getLoc(),

    CHECK(op_def != nullptr);  // Crash OK

    const ApiDef* api_def = api_def_map_->GetApiDef(op_name);
    CHECK(api_def != nullptr);  // Crash OK

    operators_.push_back(OpSpec::Create(*op_def, *api_def));
  }
}

}  // namespace generator