"args_5"}, %arg6: tensor<?x!tf_type.string> {tf._user_specified_name = "args_6"}, %arg7: tensor<?x?x512xf32> {tf._user_specified_name = "args_7"}, %arg8: tensor<?x?xf32> {tf._user_specified_name = "args_8"}, %arg9: tensor<?x300xi32> {tf._user_specified_name = "args_9"}, %arg10: tensor<?x300xi32> {tf._user_specified_name = "args_10"}, %arg11: tensor<?x300xf32> {tf._user_specified_name = "args_11"}, %arg12: tensor<?x300xf32> {tf._user_specified_name = "args_12"}, %arg13: tensor<?x!tf_type.string> {tf._user_specified_name...

// CHECK-LABEL: func @func_with_call
func.func private @some_func(%arg0: tensor<100xf32>) -> tensor<100xf32> attributes {tf.api_implements = "lstm_b4e9f0e7-ac55-42bc-8ef2-8496419a608c"}
func.func @func_with_call(%arg0: tensor<100xf32>) -> tensor<100xf32> {
  %0 = func.call @some_func(%arg0) : (tensor<100xf32>) -> tensor<100xf32>
  func.return %0 : tensor<100xf32>
  }
}

// -----

module {

      tf_device.return %5#0, %5#1, %3 : tensor<?x300xi32>, tensor<?xi32>, tensor<?x4x!tf_type.string>
    }) {device_assignment = [], num_cores_per_replica = 1 : i64, topology = ""} : () -> (tensor<?x300xi32>, tensor<?xi32>, tensor<?x4x!tf_type.string>)
    return %0#0, %0#1, %0#2 : tensor<?x300xi32>, tensor<?xi32>, tensor<?x4x!tf_type.string>
  }
}

  %r3 = "tf.Reshape"(%arg2, %shape1) : (tensor<10000xf32>, tensor<2xi32>) -> tensor<100x100xf32>
  %shape3 = arith.constant dense<[-1, 100]> : tensor<2xi32>
  %r4 = "tf.Reshape"(%arg2, %shape3) : (tensor<10000xf32>, tensor<2xi32>) -> tensor<100x100xf32>
  %r5 = "tf.Reshape"(%arg0, %arg3) : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
  %r6 = "tf.Reshape"(%arg2, %arg3) : (tensor<10000xf32>, tensor<*xi32>) -> tensor<*xf32>

      tf_device.return %5 : tensor<512xi32>
    }) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<512xi32>, tensor<1024xi32>, tensor<1024xi32>, tensor<1024xf32>)
    // CHECK: %[[COMPILE_OUTPUT:[0-9]*]]:4 = "tf_device.launch"
    // CHECK: "tf._TPUCompileMlir"()
    // CHECK: is_bounded_dynamic_dim: true

// CHECK-LABEL: testMultinomialInt32
func.func @testMultinomialInt32(%arg0: tensor<2xf32>, %arg1: tensor<1xi32>) -> tensor<10xi32> {
  // CHECK: "tfl.multinomial"(%arg0, %arg1)
  %0 = "tfl.multinomial"(%arg0, %arg1) {seed = 0 : i64, seed2 = 0: i64} : (tensor<2xf32>, tensor<1xi32>) -> tensor<10xi32>
  func.return %0 :  tensor<10xi32>
}

// -----

// CHECK-LABEL: testGelu