builder.setInsertionPoint(op);
    for (auto& op_operand : op->getOpOperands()) {
      if (IsLargeFloatType(op_operand.get().getType())) {
        op_operand.set(builder.create<TF::CastOp>(
            op->getLoc(), ToBfloat16Type(op_operand.get().getType()),
            op_operand.get()));
      }
    }
    builder.setInsertionPointAfter(op);

  auto* ctx = func.getContext();
  patterns.add<PushDownDequantize>(ctx);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    signalPassFailure();
  }
}
}  // namespace

// Creates an instance of the TensorFlow Lite optimize op order pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizeOpOrderPass() {
  return std::make_unique<OptimizeOpOrderPass>();
}

  common_error_message_.clear();
  error_collector_->Clear();
}

void ErrorCollectorInstrumentation::runAfterPassFailed(Pass *pass,
                                                       Operation *module) {
  // Create a new error if no errors collected yet.
  if (error_collector_->CollectedErrors().empty() &&
      !common_error_message_.empty()) {
    error_collector_->ReportError(NewConverterErrorData(

  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

  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()));

        }
      }

      // 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) {

    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

    StringAttr callee_sym = cluster_func_op.getFuncAttr().getAttr();
    MoveResourceArgsToEnd(symtab.lookup<func::FuncOp>(callee_sym));
  }
  builder.setInsertionPoint(cluster_func_op);
  auto xla_launch_op = builder.create<TF::XlaLaunchOp>(
      cluster_func_op.getLoc(), cluster_func_op.getResultTypes(),
      /*constants=*/ValueRange({}), ValueRange(non_resource_args),
      ValueRange(resource_args), cluster_func_op.getFuncAttr());