def LiftMatMul : Pat<
  (TF_MatMulOp:$res $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
  (LiftAsTFPartitionedCall<"composite_matmul_fn">
    (ArgumentList $a, $b),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"transpose_a"> $transpose_a),
      (NamedAttr<"transpose_b"> $transpose_b))),
  [(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;

def LiftConv3D : Pat<

  // CHECK: tf.MatMul
  %outputs = "tf.MatMul"(%arg0, %arg1) {device = "TPU", transpose_a = false, transpose_b = false} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
  %outputs_0 = "tf.MatMul"(%arg0, %arg1) {device = "TPU", transpose_a = false, transpose_b = false} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
  // CHECK-NOT: tf.IdentityN

  // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32>

    (TF_SubOp (TF_CastOp $input, $truncate), $input_zp),
    (TF_CastOp (TF_IdentityOp $weight), $truncate1),
    $transpose_a, $transpose_b, $grad_a, $grad_b),
  (CreateXlaDotV2OpFromTfMatMulOp
    $input, $weight, $input_zp,
    /*weight_zp=*/(CreateScalarIntegerConst<"int32_t", "0">), $matmul,
    $transpose_a, $transpose_b),
  [(IsInt8ElementType $input),
   (IsInt8ElementType $weight),
   (IsConstTensor $input_zp),

    return %0 : tensor<2x2xf32>
  }
  func.func private @composite_matmul_fn_2_0(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> attributes {tf_quant.composite_function} {

  %0 = corert.get_op_handler %arg0 "cpu"
  // CHECK: %[[RESULT:.*]] = corert.executeop(%[[ARG_0:.*]]) "tf.MatMul"(%[[ARG_1:.*]], %[[ARG_1]]) {T = f32, transpose_a = false, transpose_b = false} : 1
  %1 = corert.executeop(%0) "tf.MatMul"(%arg1, %arg1) {T = f32, device = "cpu", transpose_a = false, transpose_b = false} : 1
  tfrt.return %arg0, %1 : !tfrt.chain, !corert.tensorhandle

def LiftMatMul : Pat<
  (TF_MatMulOp:$res $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
  (LiftAsTFPartitionedCall<"composite_matmul_fn">
    (ArgumentList $a, $b),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"transpose_a"> $transpose_a),
      (NamedAttr<"transpose_b"> $transpose_b))),
  [(IsNotInLiftedFunc $res), (IsConstTensor $b)], [], (addBenefit 1)>;

disable this feature as well.

### Fuse SpaceToDepth with Automatic Double Transpose

The transpose and reshape op in SpaceToDepthOp on TPU hosts may cause image
model to be infeed bound. To reduce host time, space to depth transform can be
fused with `automatic double transpose` to reduce extra overhead on the host.

### Extend from Conv2D to Conv3D

        return failure();
      }
    }

    // Input rhs must be a constant with rank 2.
    if (constant.getType().getRank() != 2) return failure();

    // Create a tfl.transpose op that performs ZX transpose on `input`.
    auto create_z_x_transpose_op = [&](Value input) -> Value {
      RankedTensorType input_type =
          mlir::cast<RankedTensorType>(input.getType());

    %0 = "tf.Identity"(%cst) : (tensor<2x1024xi8>) -> tensor<2x1024xi8>
    %1 = "tf.Cast"(%0) {Truncate = false} : (tensor<2x1024xi8>) -> tensor<2x1024xf32>
    %2 = "tf.MatMul"(%arg0, %1) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    %3 = "tf.Mul"(%2, %cst_0) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
    return %3 : tensor<*xf32>
  }