%19 = "tf.Cast"(%18) {Truncate = false, device = ""} : (tensor<2x4x3x6xi32>) -> tensor<2x4x3x6xf32>
    %20 = "tf.Mul"(%19, %cst_0) {device = ""} : (tensor<2x4x3x6xf32>, tensor<f32>) -> tensor<2x4x3x6xf32>
    %21 = "tf.Relu"(%20) {device = ""} : (tensor<2x4x3x6xf32>) -> tensor<2x4x3x6xf32>
    %22 = "tf.Minimum"(%21, %cst) {device = ""} : (tensor<2x4x3x6xf32>, tensor<f32>) -> tensor<2x4x3x6xf32>

  attributes  {tf.entry_function = {inputs = "args_0,args_1", outputs = "rets_0,rets_1"}} {
    %0:2 = tf_executor.graph {
      %1:3 = tf_executor.island wraps "tf.IdentityN"(%arg0, %arg1) {T = ["tfdtype$DT_FLOAT", "tfdtype$DT_INT32"], device = "", name = "identity"} : (tensor<*xf32>, tensor<2x4x6x8xi32>) -> (tensor<*xf32>, tensor<2x4x6x8xi32>)

  attributes  {tf.entry_function = {inputs = "args_0,args_1", outputs = "rets_0,rets_1"}} {
    %0:2 = tf_executor.graph {
      %1:3 = tf_executor.island wraps "tf.IdentityN"(%arg0, %arg1) {T = ["tfdtype$DT_FLOAT", "tfdtype$DT_INT32"], device = "", name = "identity"} : (tensor<*xf32>, tensor<2x4x6x8xi32>) -> (tensor<*xf32>, tensor<2x4x6x8xi32>)

// -----

// CHECK-LABEL: doubleTposeInputPermNotEqualNoChange
func.func @doubleTposeInputPermNotEqualNoChange(%arg0: tensor<2x4x3x5xf32>, %arg1: tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32> {
  %perm = arith.constant dense<[3, 0, 2, 1]> : tensor<4xi32>
  %0 = "tfl.transpose"(%arg0, %perm) : (tensor<2x4x3x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
  %perm1 = arith.constant dense<[3, 0, 1, 2]> : tensor<4xi32>

# Verify arg and ret devices are added as arg and ret attributes.

# CHECK-LABEL: func @main
# CHECK-SAME:  (%[[ARG_0:[a-z0-9]+]]: tensor<*xf32> {tf.device = "/CPU:0"}, %[[ARG_1:[a-z0-9]+]]: tensor<2x4x6x8xi32>) -> (tensor<*xf32>, tensor<*xi32> {tf.device = "/CPU:1"})

node {
  name: "args_0"
  op: "_Arg"
  device: "/CPU:0"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {

// RUN: tf-mlir-translate -mlir-to-graphdef %s -tf-graph-as-function -o - | FileCheck %s

func.func @main(%arg0: tensor<*x!tf_type.resource>, %arg1: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>>, %arg2: tensor<*xf32>, %arg3: tensor<2x4x6x8xi32>) -> (tensor<f32>, tensor<f32>)
attributes {tf.entry_function = {inputs = "args_0,args_1,args_2,args_3", outputs = "rets_0_RetVal,rets_1_RetVal"}} {
  %graph:2 = tf_executor.graph {

# functions are converted.

# CHECK:      func @main(%arg0: tensor<*x!tf_type.resource>, %arg1: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>>, %arg2: tensor<*xf32>, %arg3: tensor<2x4x6x8xi32>) -> (tensor<*xf32>, tensor<*xf32>)
# CHECK-SAME: _xla_compile_device_type = "GPU"
# CHECK-SAME: allow_soft_placement
# CHECK-SAME: control_outputs = ""
# CHECK-SAME: inputs = "args_0,args_1,args_2,args_3"

// -----

func.func @transpose_output_type_bad(%arg0 : tensor<3x4x5x6xi32>) -> tensor<3x4x5x6xi32> {
  %cst = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
  // expected-error @+1 {{expect output type tensor<3x6x4x5xi32>, got tensor<3x4x5x6xi32>}}
  %0 = "tfl.transpose"(%arg0, %cst) : (tensor<3x4x5x6xi32>, tensor<4xi32>) -> tensor<3x4x5x6xi32>
  func.return %0 : tensor<3x4x5x6xi32>
}


// -----

func.func private @dense_image_warp_invalid_input_type(%arg0: tensor<2x4x4x1xi32>, %arg1: tensor<2x4x4x2xf32>) -> tensor<2x4x4x1xf32> attributes {tf._implements = "addons:DenseImageWarp"}

// expected-warning @+1 {{Flow should be a 4D float tensor}}
func.func private @dense_image_warp_invalid_flow_type(%arg0: tensor<2x4x4x1xf32>, %arg1: tensor<2x4x4x2xi32>) -> tensor<2x4x4x1xf32> attributes {tf._implements = "addons:DenseImageWarp"}

}

func.func @gatherWithBatchDims(%arg0 : tensor<2x3x6xf32>, %arg1 : tensor<2x5xi32>) -> tensor<2x5x3x6xf32> {
  %0 = "tf.Const"() { value = dense<[1]> : tensor<1xi32> } : () -> tensor<1xi32>
  %1 = "tf.GatherV2"(%arg0, %arg1, %0) {batch_dims = 1 : i64} : (tensor<2x3x6xf32>, tensor<2x5xi32>, tensor<1xi32>) -> tensor<2x5x3x6xf32>
  func.return %1 : tensor<2x5x3x6xf32>

// CHECK-LABEL:gatherWithBatchDims