// after this pass is run.
namespace mlir {
namespace TFL {
namespace {
#define GEN_PASS_DEF_TRIMFUNCTIONSPASS
#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

// The pass to trim functions before we legalize to TFL
// dialect using the specified allowlist.
class TrimFunctionsPass
    : public impl::TrimFunctionsPassBase<TrimFunctionsPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TrimFunctionsPass)

  %0 = stablehlo.constant dense<[[1.4, 0.0], [3.1, -0.2]]> : tensor<2x2xf32>
  return %0 : tensor<2x2xf32>
}

// -----

func.func @constant_elided() -> tensor<2x2xf32> {
  // expected-error @+1 {{failed to legalize operation 'stablehlo.constant' that was explicitly marked illegal}}
  %0 = stablehlo.constant dense_resource<__elided__> : tensor<2x2xf32>
  return %0 : tensor<2x2xf32>
}

// -----

    func.return
  }
})";

// MLIR which should not legalize at all
static constexpr char kBadMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  func.func @main() -> () {
    %0 = tf.Unknown() -> ()
    func.return %0
  }
})";

// MLIR which should be filtered by the MLIR bridge but fully legalize with the
// combined bridge.

    ],
    deps = [
        "@llvm-project//mlir:AffineDialect",
        "@llvm-project//mlir:QuantOps",
        # Link jit lib to link JIT devices required to run
        # xla-legalize-tf-with-tf2xla pass.
        "//tensorflow/compiler/jit",
        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
        "//tensorflow/compiler/mlir/lite:tensorflow_lite_legalize_tf",

      composite->getLoc(), composite->getResultTypes(),
      composite->getOperands(), composite.getName(),
      CustomOption(&builder, custom_option_buffer));
}

}  // namespace

// Legalize stablehlo::CompositeOp to TFL::CustomOp for runtime-supported
// composites. See `IsSupportedComposite` for list of supported ops.
//
// Example:
//   %0 = stablehlo.composite "odml.some_op" <args> {

constexpr char kTfLiteInputIndices[] = "_tflite_input_indices";

// Legalize operations in functions.
class LegalizeTFPass : public impl::LegalizeTFPassBase<LegalizeTFPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeTFPass)
  LegalizeTFPass() = default;
  explicit LegalizeTFPass(bool run_tfl_runtime_verification,
                          bool preserve_assert_op) {

        return WalkResult::interrupt();
      }

      return WalkResult::advance();
    });

    if (walk_result.wasInterrupted()) {
      return tsl::errors::Internal("Could not legalize all ops");
    }

    return absl::OkStatus();
  }

  mlir::func::FuncOp GetMainFunc() {
    func::FuncOp main_func = module_->lookupSymbol<mlir::func::FuncOp>("main");
    EXPECT_TRUE(main_func);

    "post-training-quantization",
    llvm::cl::desc("Enable post_training_quantization."),
    llvm::cl::init(false));

// NOLINTNEXTLINE
opt<bool> legalize_custom_tensor_list_ops(
    "legalize-custom-tensor-list-ops",
    llvm::cl::desc("Convert \"tf.TensorList*\" ops to \"tfl.custom_op\""
                   "if they can all be supported."),
    llvm::cl::init(false));

// NOLINTNEXTLINE

 public:
  explicit TflToStablehloPass() : PassWrapper() {}
  StringRef getArgument() const final { return "tfl-parse-stablehlo-ops"; }
  StringRef getDescription() const final {
    return "This pass will legalize TFLite custom Ops to StableHLO ops.";
  }

 private:
  void runOnOperation() override;

  void getDependentDialects(DialectRegistry& registry) const override {

                         OperationPass<ModuleOp>> {
 public:
  StringRef getArgument() const final {
    return "tf-xla-call-module-op-to-stablehlo-pass";
  }
  StringRef getDescription() const final {
    return "Legalize TF_XlaCallModule Op to stablehlo";
  }
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<stablehlo::StablehloDialect, vhlo::VhloDialect,