opt<bool> unfold_batchmatmul(
    "unfold_batchmatmul",
    llvm::cl::desc(
        "Whether to unfold TF BatchMatMul to a set of TFL FullyConnected ops."),
    llvm::cl::init(false));

// NOLINTNEXTLINE
opt<bool> unfold_large_splat_constant(
    "unfold-large-splat-constant",
    llvm::cl::desc("Whether to unfold large splat constant tensors to reduce "
                   "the generated model size."),

// RUN: odml-to-stablehlo-opt %s -unfold-splat-constant-pass -cse -verify-diagnostics | FileCheck %s

// CHECK-LABEL: @unfold_splat_constant_float
func.func @unfold_splat_constant_float() -> tensor<1x750xf32> {
  %cst = mhlo.constant dense<7.680000e+02> : tensor<1x750xf32>
  func.return %cst : tensor<1x750xf32>

  // CHECK-DAG: %0 = mhlo.constant dense<7.680000e+02> : tensor<f32>

#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {
namespace {

#define DEBUG_TYPE "unfold-splat-constant-pass"

#define GEN_PASS_DEF_UNFOLDSPLATCONSTANTPASS
#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h.inc"

// Undo the MHLO::BroadcastInDimOp folding pattern on splat tensor.

  let summary = "Legalize StableHLO custom call ops where the call target is 'stablehlo.composite' to composite ops.";
  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
}

def UnfoldSplatConstantPass : Pass<"unfold-splat-constant-pass", "ModuleOp"> {
  let summary = "Replaces a splat constant tensor with a BroadcastInDim op.";
  let constructor = "mlir::odml::CreateUnfoldSplatConstantPass()";
}

  // By default we freeze all variables and disallow mutable variables. When
  // 'enable_tflite_variables' is true then we allow mutable variable only.
  bool disable_variable_freezing = false;
  // Whether to unfold large splat constant tensors and replace them with
  // fill operation.
  bool unfold_large_splat_constant = false;
  // Whether to run the `GuaranteeAllFuncsOneUsePass` to ensure each function
  // has a single use.