// RUN: tf-tfrt-opt -tf-to-tfrt %s | FileCheck %s --dump-input=fail

// CHECK-LABEL: func @device_test
func.func @device_test(
    %arg0: tensor<3x1xf32> {tf_saved_model.index_path = [0]},
    %arg1: tensor<1x3xf32> {tf_saved_model.index_path = [0]})
      -> (tensor<3x3xf32> {tf_saved_model.index_path = []}) {
  // CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"

  func.return %1 : tensor<?x?xf32>
}

// -----

// CHECK-LABEL: func @uniform_quantized_add
func.func @uniform_quantized_add(%input: tensor<3x2xf32>) -> tensor<3x2xf32> {
  %input_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
  %input_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>

func.func @device_test(%arg0: tensor<3x1xf32>) -> (tensor<3x3xf32>) {

  // CHECK: device = "gpu"
  %0 = "tf.Const"() {value = dense<[[1.0, 2.0, 3.0]]> : tensor<1x3xf32>} : () -> tensor<1x3xf32>
  // CHECK: device = "gpu"
  %1 = "tf.MatMul"(%arg0, %0) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>

  constexpr char mlir_module_string[] = R"mlir(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main(%arg0 : tensor<2x2xf32>) -> tensor<2x2xf32> {
      %max = "tf.Const"() { value = dense<12.0> : tensor<f32> } : () -> tensor<f32>
      %min = "tf.Const"() { value = dense<-25.0> : tensor<f32> } : () -> tensor<f32>

// CHECK-NEXT:    return %2 : tensor<3x2xf32>
}

func.func @transpose_3d(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> {
  %0 = "mhlo.transpose"(%arg0) <{permutation = dense<[2, 1, 0]> : tensor<3xi64>}> : (tensor<1x2x3xf32>) -> tensor<3x2x1xf32>
  func.return %0 : tensor<3x2x1xf32>

// CHECK-LABEL:   transpose_3d

  // CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
  %cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32>
  // CHECK: %[[q2:.*]] = "quantfork.qcast"(%[[ARG_0]])
  // CHECK-SAME: quant.uniform<i4:f32, 0.13329366842905679:-1>
  // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]])

  # Verify that the function defined using function.Defun
  # has a corresponding tf.LegacyCall op.
  # CHECK:      func {{@[a-zA-Z_0-9]+}}(
  # CHECK-SAME: [[ARG0:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["y"]},
  # CHECK-SAME: [[ARG1:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]}
  #
  # CHECK-NEXT: [[R0:%.*]] = "tf.LegacyCall"([[ARG1]], [[ARG0]])
  z = plus(x, y)

  func.return %0: tensor<3x2xf32>
}

func.func @sparse_f16() -> tensor<3x2xf16> {
  // CHECK-LABEL: @sparse_f16

# CHECK:      func {{@[a-zA-Z_0-9]+}}(
# CHECK-SAME:   [[ARG0:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]},
# CHECK-SAME:   [[ARG1:%.*]]: tensor<!tf_type.resource<tensor<1x3xf32>>> {tf_saved_model.bound_input = @[[VAR]]})
# CHECK-SAME:             -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]})

  module {
    func.func @constant_add() -> (tensor<3x2xf32>) {
      %cst1 = stablehlo.constant dense<2.4> : tensor<3x2xf32>
      %cst2 = stablehlo.constant dense<5.7> : tensor<3x2xf32>
      %add = stablehlo.add %cst1, %cst2 : (tensor<3x2xf32>, tensor<3x2xf32>) -> tensor<3x2xf32>
      func.return %add : tensor<3x2xf32>
    }
  }
)mlir";