// CHECK: [[NONE:%.*]] = "tfl.no_value"() <{value}> : () -> none
// CHECK: "tfl.fully_connected"(%arg0, %arg1, [[NONE]])
// CHECK-SAME: (tensor<40x37xf32>, tensor<40x37xf32>, none) -> (tensor<40x40xf32>, tensor<40x40xf32>)
  %cst = "tfl.no_value"() {value = unit} : () -> none

  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tfl.transpose"(%arg1, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>
  // CHECK: "tfl.batch_matmul"(%arg0, %arg1)
  %1 = "tfl.fully_connected"(%arg0, %0, %arg2) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32>

  %cst_0 = arith.constant dense<2.0> : tensor<256x10368xbf16>
  %cst_1 = "tfl.no_value"() {value = unit} : () -> none
  %cst_2 = arith.constant dense<3.0> : tensor<32x1x256xbf16>
  %0 = "tfl.fully_connected"(%arg0, %cst_0, %cst_1) {
    fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"
  } : (tensor<1x10368xbf16>, tensor<256x10368xbf16>, none) -> tensor<1x256xbf16>

func.func @main(tensor<40x37xf32>, tensor<40x37xf32>) -> tensor<40x40xf32> {
^bb0(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>):
  %cst = arith.constant dense<1.0> : tensor<40xf32>
  %0:2 = "tfl.fully_connected"(%arg0, %arg1, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<40x37xf32>, tensor<40x37xf32>, tensor<40xf32>) -> (tensor<40x40xf32>, tensor<40x40xf32>)

func.func @main(tensor<40x37xf32>, tensor<40x37xf32>) -> tensor<40x40xf32> {
^bb0(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>):
  %cst = arith.constant dense<1.0> : tensor<40xf32>
  %0:2 = "tfl.fully_connected"(%arg0, %arg1, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<40x37xf32>, tensor<40x37xf32>, tensor<40xf32>) -> (tensor<40x40xf32>, tensor<40x40xf32>)

  // FB-NEXT:              max: [ 1.0 ]
  // FB-NEXT:      }

  // FB-LABEL:     name: "arg1",
  // FB-NEXT:            quantization: {
  // FB-EMPTY:
  // FB-NEXT:            }

  // FB-LABEL:     name: "tfl.fully_connected",
  // FB-NEXT:      quantization: {
  // FB-NEXT:        min: [ -0.0, -1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0,
  // FB-SAME:  -8.0, -9.0, -10.0, -11.0, -12.0, -13.0, -14.0, -15.0, -16.0,

# CHECK:           %[[VAL_10:.*]] = "tfl.fully_connected"(%[[VAL_7]]#0, %[[VAL_9]], %[[VAL_3]]) <{fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x5x3xf32>, tensor<7x3xf32>, none) -> tensor<5x7xf32>

        return false;
      }
      total_count += input_type.getNumElements();
    }
    *count = total_count;
    return true;
  }

  // For conv2d/depthwise_conv/fully_connected ops.
  // This algorithm actually comes from TOCO tooling_util.cc
  static bool GetArithmeticCountForConvAndFullyconnectedOp(mlir::Operation* op,
                                                           int64_t* count) {

  OptimizeBatchMatmulPass() = default;
  OptimizeBatchMatmulPass(const OptimizeBatchMatmulPass &) {}

  void runOnOperation() override;
};

// Converts batch_matmul operation to fully_connected if rhs is a
// constant tensor with rank 2
struct ConvertBatchMatMulOp2FullyConnectedOp
    : public OpRewritePattern<TFL::BatchMatMulOp> {
  using OpRewritePattern<TFL::BatchMatMulOp>::OpRewritePattern;

  // 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.
  bool shape_inference = true;