func.return %0 : tensor<? x i32>
}

// CHECK-LABEL: testPow
func.func @testPow(tensor<? x i32>, tensor<? x i32>) -> tensor<? x i32> {
^bb0(%arg0: tensor<? x i32>, %arg1: tensor<? x i32>):
  // CHECK: tfl.pow %arg0, %arg1
  %0 = tfl.pow %arg0, %arg1 : tensor<? x i32>
  func.return %0#0 : tensor<? x i32>
}

// CHECK-LABEL: testAtan2

    %0 = "tf.D"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
    %1 = func.call @referenced_func2(%0) : (tensor<?xi32>) -> tensor<?xi32>
    func.return %1 : tensor<?xi32>
  }

  func.func @referenced_func2(%arg0: tensor<?xi32>) -> tensor<?xi32> {
    %0 = "tf.E"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
    func.return %0 : tensor<?xi32>
  }
}

// -----

    // CHECK: tf.BroadcastGradientArgs
    // CHECK-SAME: (tensor<*xi32>, tensor<0xi32>) -> (tensor<?xi32>, tensor<?xi32>)
    %2:2 = "tf.BroadcastGradientArgs"(%arg0, %arg1) {T = "tfdtype$DT_INT32", name = "BroadcastGradientArgs"} : (tensor<*xi32>, tensor<0xi32>) -> (tensor<?xi32>, tensor<?xi32>)
    func.return %2#0 : tensor<?xi32>
  }

  let arguments = (ins
     TFL_TensorOf<[I8, I16, I32, I64, UI8, UI16, UI32, UI64, F32, F64]>:$input,
     TFL_TensorOf<[I32]>:$dilations,
     TFL_0DTensorOf<[I8, I16, I32, I64, UI8, UI16, UI32, UI64, F32, F64]>:$padding_value
  );

  let results = (outs TFL_TensorOf<[I8, I16, I32, I64, UI8, UI16, UI32, UI64, F32, F64]>:$output);
}

def TFL_AddOp : TFL_Op<"add", [

    ```mlir
      %2 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
      %3 = "tf.B"(%2) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>
      %4 = "tf.C"(%2, %3) {device = "tpu0"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
      %5 = "tf.D"(%4) : (tensor<?xi32>) -> tensor<?xi32>
    ```

    After the pass, we will have:

    ```mlir
      %0 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>

    "tf_device.cluster"() ({
      %0 = "tf.A"() : () -> (tensor<i32>)
      %1 = "tf.B"() {_xla_outside_compilation = "cluster0"} : () -> (tensor<i32>)
      %2 = "tf.C"(%1) : (tensor<i32>) -> (tensor<i32>)
      %3 = "tf.D"() : () -> (tensor<i32>)
      "tf.E"(%2, %3) : (tensor<i32>, tensor<i32>) -> ()
      "tf.F"(%1, %0) {_xla_outside_compilation = "cluster0"} : (tensor<i32>, tensor<i32>) -> ()
      "tf.G"(%0, %1) : (tensor<i32>, tensor<i32>) -> ()

  // expected-error @+1 {{argument count mismatch: 'args' has 1 arguments, but '@pcall_func_2' expects 2}}
  %0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @pcall_func_2} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}

func.func @pcall_func_2(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
  func.return %arg0 : tensor<i32>
}

// -----

    },
    {
      // loop body
      ^bb0(%barg0: tensor<i32>, %barg1: tensor<i32>):
      %one = arith.constant dense<1> : tensor<i32>
      %sub0 = "tf.Sub"(%barg0, %one) : (tensor<i32>, tensor<i32>) -> tensor<i32>
      %sub1 = "tf.Sub"(%barg1, %one) : (tensor<i32>, tensor<i32>) -> tensor<i32>
      "tf.Yield"(%sub0, %sub1) : (tensor<i32>, tensor<i32>) -> ()
    }

}

func.func @argmin(%arg0: tensor<3xi32>, %arg1: tensor<i32>) -> tensor<i32> {
  %0 = "tf.ArgMin"(%arg0, %arg1) : (tensor<3xi32>, tensor<i32>) -> tensor<i32>
  func.return %0 : tensor<i32>

// CHECK-LABEL: argmin
// CHECK:  "tfl.arg_min"(%arg0, %arg1) : (tensor<3xi32>, tensor<i32>) -> tensor<i32>
}

func.func @sigmoid(%arg0: tensor<?x88xf32>) -> tensor<?x88xf32> {

func.func @embedding(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) -> tensor<*xf32> attributes  {tf._implements = "embedding_matmul", tf._reference = "mlir"} {
  %0 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
  %1 = "tf.ExpandDims"(%arg1, %0) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
  %2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
  %3 = "tf.Const"() {value = dense<4096> : tensor<i32>} : () -> tensor<i32>