# RUN: tf_tfl_translate -unfold_batchmatmul=false -tf-input-arrays=Placeholder,Placeholder_1 -tf-input-shapes=2,5,3:3,7 -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-output-arrays=MatMul -output-mlir %s -o - 2>&1 | FileCheck %s

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2

# RUN: tf_tfl_translate -tf-input-arrays=Placeholder,Placeholder_1 -tf-input-shapes=2,5,3:3,7 -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-output-arrays=MatMul -unfold_batchmatmul=true -output-mlir %s -o - 2>&1 | FileCheck %s

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2
        }

  mlir::TFL::PassConfig pass_config(quant_specs);
  bool emit_builtin_tflite_ops = !toco_flags.force_select_tf_ops();
  pass_config.emit_builtin_tflite_ops = emit_builtin_tflite_ops;
  pass_config.unfold_batch_matmul = toco_flags.unfold_batchmatmul();
  pass_config.lower_tensor_list_ops = toco_flags.lower_tensor_list_ops();
  pass_config.legalize_custom_tensor_list_ops =
      toco_flags.legalize_custom_tensor_list_ops();

    "select-user-tf-ops",
    llvm::cl::desc(
        "<list of custom tf ops created by the user (comma separated)>"),
    llvm::cl::init(""));

// NOLINTNEXTLINE
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(

extern llvm::cl::opt<bool> convert_tf_while_to_tfl_while;
extern llvm::cl::opt<std::string> select_user_tf_ops;
extern llvm::cl::opt<bool> allow_all_select_tf_ops;
extern llvm::cl::opt<bool> unfold_batchmatmul;
extern llvm::cl::opt<bool> unfold_large_splat_constant;
extern llvm::cl::opt<bool> guarantee_all_funcs_one_use;
extern llvm::cl::opt<bool> enable_dynamic_update_slice;
extern llvm::cl::opt<bool> preserve_assert_op;

  pass_config.unfold_batch_matmul = toco_flags.unfold_batchmatmul();
  pass_config.lower_tensor_list_ops = toco_flags.lower_tensor_list_ops();
  // Disable the unfolding of the 16x16 TF::BatchMatMulOp to avoid the
  // conversion to an unsupported 16x16 TFL::FullyConnectedOp.
  if (toco_flags.inference_type() == toco::IODataType::QUANTIZED_INT16) {
    pass_config.unfold_batch_matmul = false;
  }

  pass_config.unfold_batch_matmul = toco_flags.unfold_batchmatmul();
  pass_config.lower_tensor_list_ops = toco_flags.lower_tensor_list_ops();
  // Disable the unfolding of the 16x16 TF::BatchMatMulOp to avoid the
  // conversion to an unsupported 16x16 TFL::FullyConnectedOp.
  if (toco_flags.inference_type() == toco::IODataType::QUANTIZED_INT16) {
    pass_config.unfold_batch_matmul = false;
  }

  }

  mlir::TFL::PassConfig pass_config(quant_specs);
  pass_config.emit_builtin_tflite_ops = emit_builtin_tflite_ops;
  pass_config.lower_tensor_list_ops = lower_tensor_list_ops;
  pass_config.unfold_batch_matmul = unfold_batchmatmul;
  pass_config.unfold_large_splat_constant = unfold_large_splat_constant;
  pass_config.guarantee_all_funcs_one_use = guarantee_all_funcs_one_use;

    "mlir::quant::QuantizationDialect",
    "mlir::quantfork::QuantizationForkDialect",
    "mhlo::MhloDialect"
  ];
  let options = [
      Option<"unfold_batch_matmul_", "unfold_batchmatmul",
             "bool", "true",
             "Unfold BatchMatMul into individual MatMul ops.">,
      Option<"allow_bf16_and_f16_type_legalization_", "allow-bf16-and-f16-type-legalization",

  // ops of the same device, under a `tf_device.launch` op.
  bool form_clusters = false;
  // If `unfold_batch_matmul` is true, the tf.BatchMatMul is unfolded to a set
  // of tfl.fully_connected ops.
  bool unfold_batch_matmul = true;
  // Whether to outline WhileOp at the end of the pipeline.
  bool outline_tf_while = false;
  // Whether to do shape inference.