// CHECK: return %[[FULLY_CONNECTED]]

// -----

// Tests that when the weight tensor for `stablehlo.dot_general` has a
// `stablehlo.constant` -> `stablehlo.transpose` pattern, the
// `stablehlo.constant` is directly transformed to `tfl.pseudo_qconst`, which
// becomes the rhs of `tfl.fully_connected`. This is because
// `tfl.fully_connected` accepts a [o, i] format for rhs, which

// filter value from [i, o] -> [o, i]. This is because we assume `[i, o]`
// format for `stablehlo.dot_general` (i.e. contracting dimension == 1)
// whereas `tfl.fully_connected` accepts an OI format.
TFL::QConstOp CreateTransposedTflConstOpForFilter(
    stablehlo::ConstantOp filter_constant_op, PatternRewriter& rewriter,
    bool is_per_channel) {

// Remove Reshape before FullyConnected when `keep_num_dims=false` and Reshape
// does not alter the last dimension as FullyConnected will collapse all other
// dimensions into a single dimension. For example,
//
//   %shape = arith.constant dense<[1, 128, 64]> : tensor<3xi32>
//   %reshape = tfl.reshape(%input, %shape) // %input: tensor<128x64xf32>
//   %fc = tfl.fully_connected(%reshape, %filter, %bias)

  func.return %0, %1, %2 : tensor<2xi32>, tensor<2xf32>, tensor<i32>
}

// -----

func.func @fully_connected(%arg0: tensor<1x37xf32>, %arg1: tensor<40x37xf32>, %arg2: tensor<40xf32>) -> tensor<1x40xf32> {

// CHECK: %[[ARG:.*]] = "tfl.transpose"(%arg1, %[[CST]]) : (tensor<37x40xf32>, tensor<2xi32>) -> tensor<40x37xf32>
// CHECK: %[[CST_0:.*]] = "tfl.no_value"() <{value}> : () -> none
// CHECK: "tfl.fully_connected"(%arg0, %[[ARG]], %[[CST_0]]) <{fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<40x37xf32>, tensor<40x37xf32>, none) -> tensor<40x40xf32>
}

    std::vector<int> GetQuantizableOperandIndices() { return {1}; }
  }];
}

// TODO(jpienaar): Update post discussion on semantics of FC OP.
def TFL_FullyConnectedOp : TFL_Op<"fully_connected", [
    Pure, AccumulatorUniformScale<2, 0, 1>,
    AffineQuantizedOpInterface,
    AffineOpCoefficient<0, 1>,
    TFL_SparseOp,
    DeclareOpInterfaceMethods<TFL_ArithmeticCount>,
    QuantizableResult,