func.func @testConv(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>) -> tensor<256x30x30x16xf32> {
  // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"

// -----

func.func @multiple_return_no_controls_needed(%arg0: tensor<*xi32>, %arg1: tensor<i32>) -> (tensor<*xi32>, tensor<*xi32>) {
  %graph:2 = tf_executor.graph {
    %add1, %add1_control = tf_executor.island wraps "tf.Add"(%arg0, %arg1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
    %add2, %add2_control = tf_executor.island wraps "tf.Add"(%add1, %arg1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>

    tf_executor.fetch %arg0 : tensor<*xf32>
  }
  func.return %0 : tensor<*xf32>
}

// -----

func.func @invalid_nextiteration(%arg0: tensor<*xf32>, %arg1: !tf_executor.token) -> tensor<*xf32> {
  %0 = tf_executor.graph {
    "tf_executor.NextIteration.Sink"(%arg1, %arg0) : (!tf_executor.token, tensor<*xf32>) -> ()

    return %arg0 : tensor<i1>
  }
  func.func private @while_cond(%arg0: tensor<i1>) -> tensor<i1> {
    return %arg0 : tensor<i1>
  }
}

// -----
// A test verifying the replication region of TPUReplicateMetadataOp ops. Same logic tests too many TPUCompilationResultOp ops.
module {

		{name: "FMSUB", argLength: 3, reg: fp31, asm: "FMSUB"},   // arg0*arg1 - arg2
		{name: "FMSUBS", argLength: 3, reg: fp31, asm: "FMSUBS"}, // arg0*arg1 - arg2

		{name: "SRAD", argLength: 2, reg: gp21cxer, asm: "SRAD"}, // signed arg0 >> (arg1&127), 64 bit width (note: 127, not 63!)
		{name: "SRAW", argLength: 2, reg: gp21cxer, asm: "SRAW"}, // signed arg0 >> (arg1&63), 32 bit width

// -----

module {
func.func @inference_can_fuse(%arg0: tensor<?x8x8xf32>, %arg1: tensor<8x10xf32>, %arg2: tensor<8x10xf32>, %arg3: tensor<8x40xf32>, %arg4: tensor<10x40xf32>, %arg5: tensor<40xf32>) {
  %0 = "tf.Const"() {_output_shapes = ["tfshape$"], device = "", dtype = f32, value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>

}

# CHECK-NO-SUBTYPE: func @main(%arg0: tensor<i32>, %arg1: tensor<!tf_type.variant>, %arg2: tensor<i8>, %arg3: tensor<f64>, %arg4: tensor<!tf_type.resource>) -> (tensor<i32>, tensor<!tf_type.variant>, tensor<i8>, tensor<f64>, tensor<!tf_type.resource>)

// CHECK-LABEL: func @testIdentityN
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>, [[ARG1:%.*]]: tensor<f32>)
func.func @testIdentityN(%arg0: tensor<i32>, %arg1: tensor<f32>) {
  // CHECK-NOT:  "tf.IdentityN"
  // CHECK:      "tf_device.cluster"
  // CHECK-NEXT: tf_device.return [[ARG0]], [[ARG1]]
  %0:2 = "tf_device.cluster"() ({
    %1:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>)

func.func @test_batch_matmul_to_einsum(%arg0: tensor<1x2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x2x4xf32> {
  // CHECK-LABEL: test_batch_matmul_to_einsum
  // CHECK: "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...kn->...mn"}> : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32>
  %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32>

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

func.func @einsum_broadcast4(%arg0: tensor<3x4x5x6x7xf32>, %arg1: tensor<7x8xf32>) -> tensor<3x4x5x6x8xf32> {