# 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>

// CHECK: {
// CHECK-NEXT:  version: 3,
// CHECK-NEXT:  operator_codes: [ {
// CHECK-NEXT:    deprecated_builtin_code: 9,
// CHECK-NEXT:    version: 1,
// CHECK-NEXT:    builtin_code: FULLY_CONNECTED
// CHECK-NEXT:  } ],
// CHECK-NEXT:  subgraphs: [ {
// CHECK-NEXT:    tensors: [ {
// CHECK-NEXT:      shape: [ 1, 3 ],
// CHECK-NEXT:      buffer: 1,
// CHECK-NEXT:      name: "serving_default_input2:0",

func.func @main(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>) -> tensor<40x40xf32> {
  %0 = "tfl.no_value"() {value = unit} : () -> none
  %1:2 = "tfl.fully_connected"(%arg0, %arg1, %0) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<40x37xf32>, tensor<40x37xf32>, none) -> (tensor<40x40xf32>, tensor<40x40xf32>)
  func.return %1 : tensor<40x40xf32>
}

  func.return %0 : tensor<1x112x112x32xf32>
}

func.func @fully_connected(%arg0: tensor<1x37xf32>, %arg1: tensor<40x37xf32>, %arg2: tensor<40xf32>) -> tensor<1x40xf32> {
  // CHECK: _arithmetic_count = 3000 : i64

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

  %0 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>

  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;

  TargetHardwareOpRegistration<CpuHardware, Op> Op##_CpuHardware_hardware( \
      Create);

// Operation costs on CPU

// Currently used for these ops:
// tfl.conv_2d / tfl.depthwise_conv_2d / tfl.fully_connected
class CpuConvOp : public TargetHardwareOperation {
  double GetOpCost(mlir::Operation* op) const override {
    float cost = 0.0;
    int64_t arithmetic_count;

      (ins "bool":$sign, "int":$bit_width)
    >,
  ];
}

def AffineQuantizedOpInterface : OpInterface<
  "AffineQuantizedOpInterface"> {
  let description = [{
    Interface for affine quantized ops (conv2d, fully_connected, etc.)
  }];

  let methods = [
    InterfaceMethod<
      [{Returns the affine operand index.}],
      "int", "GetAffineOperandIndex",
      (ins), [{}], [{return 1;}]>,
    InterfaceMethod<