// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<-131072> : tensor<1x3xi32>}> : () -> tensor<1x3xi32>
// CHECK: %[[MATMUL:.*]] = "tf.XlaDotV2"({{.*}}, %[[WEIGHT]])
// CHECK-SAME: (tensor<1x1024xi8>, tensor<1024x3xi8>) -> tensor<1x3xi32>
// CHECK: %[[SUB:.*]] = "tf.Sub"(%[[MATMUL]], %[[CONST]]) : (tensor<1x3xi32>, tensor<1x3xi32>) -> tensor<1x3xi32>
}

// -----

module attributes {} {

// Run optimize-batch-matmul pass only and check the results.
// RUN: tf-opt %s -tfl-optimize-batch-matmul | FileCheck %s

// CHECK-LABEL: FuseTransposeFCRhsToBatchMatmul
func.func @FuseTransposeFCRhsToBatchMatmul(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128xf32>, %arg2: none) -> tensor<16x128xf32> {
  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tfl.transpose"(%arg1, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>

                         %input : tensor<*xf32>, %weight : tensor<*x!tf_type.qint8>,
                         %weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
      attributes {tf_quant.quantized_ops = ["MatMul"]} {

    %out = "tf.UniformQuantizedDotHybrid"(%input, %weight,
                                %weight_scale, %weight_zp) {
        Tlhs = "tfdtype$DT_FLOAT",
        Trhs = "tfdtype$DT_QINT8",

  %ch1 = tfrt.merge.chains %ch, %ch0 : !tfrt.chain, !tfrt.chain

  %ch2 = tfrt_fallback_async.createop(%ch1) key(0) device("/CPU:0") "tf.MatMul"() {T = i32} num_args(2)

  %ch3, %result = tfrt_fallback_async.executeop.seq(%ch2) key(0) cost(100) device("/CPU:0") "tf.MatMul"(%a, %b) {T = i32}  : 1

  %s = "tfrt_test.get_string"() { value = "Running @matmul_cpu" } : () -> !tfrt.string

#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/util/equal_graph_def.h"

namespace tensorflow {
namespace {

using ops::Complex;
using ops::Const;
using ops::Div;
using ops::MatMul;
using ops::Placeholder;
using ops::Real;
using ops::Split;
using ops::Square;
using ops::Stack;
using ops::Sub;
using ops::Unstack;

TEST(GradientCheckerTest, BasicFloat) {

def Test():

  x = tf.constant([[1.0], [1.0], [1.0]])
  y = tf.get_variable(
      name='y',
      shape=(1, 3),
      initializer=tf.random_normal_initializer(),
      trainable=True)
  r = tf.matmul(x, y)

  tensor_info_x = tf.saved_model.utils.build_tensor_info(x)
  tensor_info_r = tf.saved_model.utils.build_tensor_info(r)

  signature_def = tf.saved_model.signature_def_utils.build_signature_def(

// NEXT ID: 7
message UnitWiseQuantizationSpec {
  // Quantization unit granularity.
  // NEXT ID: 4
  message QuantizationUnit {
    // Type of the op, ex: Conv2D, MatMul, Einsum... The node_name field can
    // be omitted if it is intended to match all nodes with this type.
    string op_type = 1;
    // Name of the node. This field accepts re2 regex format. If the node name

==============================================================================*/

// This transformation pass decomposes dense operations that assume
// support for hybrid quantization. These cases cover when a dense operation
// (e.g. matmul) has both quantized and unquantized inputs by dequantizing
// the quantized inputs, performing the operation in the expressed type, then
// requantizing if a quantized output is required.
//

  %dq_weight = "quantfork.dcast"(%q_weight) : (tensor<144x12x!quant.uniform<i8:f32, 0.074855112561992565:-1>>) -> tensor<144x12xf32>
  %9 = "tf.MatMul"(%7, %dq_weight) {transpose_a = false, transpose_b = false} : (tensor<*xf32>, tensor<144x12xf32>) -> tensor<*xf32>
  %10 = "quantfork.qcast"(%9) {volatile} : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 4.000000e-03:-12>>

             AbstractTensorHandle* const y, AbstractTensorHandle** z,
             const char* name = nullptr, const char* raw_device_name = nullptr);

// Multiply the matrix "a" by the matrix "b".
Status MatMul(AbstractContext* ctx, AbstractTensorHandle* const a,
              AbstractTensorHandle* const b, AbstractTensorHandle** product,
              bool transpose_a = false, bool transpose_b = false,