merge_fusion_with_dequantize: bool,
  ):
    lhs_dim_size, rhs_dim_size = dim_sizes
    input_shape = (*lhs_dim_size,)
    filter_shape = (*rhs_dim_size,)
    static_input_shape = [dim if dim is not None else 2 for dim in input_shape]
    model = self._create_matmul_model(
        input_shape,
        filter_shape,
        self._input_saved_model_path,
        bias_fn,
        activation_fn,
    )

    // 2. Permute dimensions of `reshaped` to produce `permuted` of shape
    //      [batch / prod(block_shape),
    //
    //       input_shape[1], block_shape[0],
    //       ...,
    //       input_shape[M], block_shape[M-1],
    //
    //       input_shape[M+1], ..., input_shape[N-1]]
    SmallVector<int64_t> permutation(reshaped_shape.size());
    permutation[0] = block_rank;
    for (int i = 0; i < block_rank; ++i) {

    int num_input_dims = ranked_input_type.getRank();
    SmallVector<int32_t, 4> padding_begin(num_input_dims, 0);
    auto input_shape = ranked_input_type.getShape();
    SmallVector<int32_t, 4> padding_end(input_shape.begin(), input_shape.end());
    SmallVector<int32_t, 4> padding_strides(num_input_dims, 1);

    int begin_mask = strided_slice_op.getBeginMask();

        size_diff, scalar_zero);

    // Build the argument/result types for if branch function.
    auto input_shape = rewriter.create<TF::ShapeOp>(
        loc, tensorflow::GetTypeFromTFTensorShape({-1}, shape_dtype),
        input_handle);

    Type branch_args_type[] = {input_handle.getType(), input_shape.getType(),
                               size_diff.getType(), size.getType()};
    Type branch_result_type[] = {result_type};

    llvm::ArrayRef<int64_t> input_shape = input_type.getShape();

    // If weight_shape[2] != 1, it means weight shape was already restored.
    if (weight_shape[2] != 1) return failure();

    // Weight was reshaped into [H, W, 1, InxMul].
    // Since we know in_channels from input_shape, we can derive multiplier.
    int64_t in_channels = input_shape[3];