// RUN: tf-opt -xla-legalize-tf %s | FileCheck %s

//===----------------------------------------------------------------------===//
// tf.BatchMatMulV2 op legalizations.
//===----------------------------------------------------------------------===//

func.func @batchmatmulv2_basic(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
// CHECK-LABEL:   func @batchmatmulv2_basic

  %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ijk,ikm->ijm"}: (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32>
  func.return %0 : tensor<3x4x6xf32>
  // CHECK-LABEL: einsum_basic
  // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32>
}

func.func @einsum_matmul(%arg0: tensor<7x9xf32>, %arg1: tensor<9x5xf32>) -> tensor<7x5xf32> {

  %0 = "tf.BatchMatMulV2"(%arg0, %arg3) {adj_x = false, adj_y = false} : (tensor<?x8x8xf32>, tensor<8x40xf32>) -> tensor<?x8x40xf32>
  %1 = "tf.Add"(%0, %arg5) : (tensor<?x8x40xf32>, tensor<40xf32>) -> tensor<?x8x40xf32>
  %2 = "tf.BatchMatMulV2"(%1, %arg4) {adj_x = false, adj_y = true} : (tensor<?x8x40xf32>, tensor<10x40xf32>) -> tensor<?x8x10xf32>

    key: "shape"
    value {
      shape {
        dim {
          size: 3
        }
        dim {
          size: 7
        }
      }
    }
  }
}
node {
  name: "MatMul"
  op: "BatchMatMulV2"
  input: "Placeholder"
  input: "Placeholder_1"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "adj_x"
    value {
      b: false
    }
  }

    key: "shape"
    value {
      shape {
        dim {
          size: 3
        }
        dim {
          size: 7
        }
      }
    }
  }
}
node {
  name: "MatMul"
  op: "BatchMatMulV2"
  input: "Placeholder"
  input: "Placeholder_1"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "adj_x"
    value {
      b: false
    }
  }

    func.return %0: tensor<2xi1>
  }

  // CHECK-LABEL: batchmatmulv2
  func.func @batchmatmulv2(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
    // CHECK: mhlo.reduce
    // CHECK: mhlo.dot_general
    // CHECK: mhlo.transpose

}

INSTANTIATE_TEST_SUITE_P(
    BatchMatMulTest, BatchMatMulTest,
    ::testing::ValuesIn<MatMulTestCase>({
        {"BatchMatMul"},
        {"BatchMatMulV2"},
        {"BatchMatMulV3"},
    }),
    [](const ::testing::TestParamInfo<BatchMatMulTest::ParamType>& info) {
      return info.param.mat_mul_method;
    });

TEST(LegalizeTFTest, DumpsProducedHLO) {

    %2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
    %3 = "tf.Sub"(%2, %weight_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>

    %5 = "tf.BatchMatMulV2"(%1, %3) {
      attr_map = "adj_x:0,adj_y:1"
    } : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
    func.return %5 : tensor<*xi32>
  }

   // Einsum with int32 accumulation.

  for (int64_t i = 0; i < input.size(); ++i) {
    output[permutation[i]] = input[i];
  }
  return output;
}

// Computes the transpositions required to convert dnums to one supported by
// tf.BatchMatmulV2 and returns the new set of dimension numbers with them.
// Transposed LHS shape will be B0,...,Bn,L0,...,Ln,C0,...,Cn and,
// transposed RHS shape will be B0,...,Bn,C0,...,Cn,R0,...,Rn respectively.

  func.return %0 : tensor<3x4x6xf32>
}
// CHECK-LABEL: lower_einsum
// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32>

// -----

func.func @removing_identity_after_const(%arg0: tensor<*xf32>) -> (tensor<*xf32>) {