// CHECK: "tf.XlaLaunch"(%arg1, %arg0) <{function = @func_with_resources, operandSegmentSizes = array<i32: 0, 1, 1>}> : (tensor<i32>, tensor<!tf_type.resource>) -> tensor<i32>
    %0 = "tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_resources} : (tensor<!tf_type.resource>, tensor<i32>) -> tensor<i32>

     %3 = builtin.unrealized_conversion_cast to tensor<i32>
     // CHECK: TPUCopyWithDynamicShape{{.*}}device = "foobar"
     %4, %5 = "tf.TPUCopyWithDynamicShape"(%0, %1, %2, %3) {operandSegmentSizes = array<i32: 2, 2>} : (tensor<2048xi32>, tensor<2048xi32>, tensor<i32>, tensor<i32>) -> (tensor<2048xi32>, tensor<2048xi32>)
    "tf.TPUExecute"(%4, %arg0) {device = "foobar"} : (tensor<2048xi32>, tensor<!tf_type.string>) -> ()

  func.return %0#0 : tensor<*xi32>
}

// CHECK: func private @compute_0

  func.return %0#0 : tensor<*xi32>
}

  func.return %1 : tensor<*xf32>
}

// -----

// The num_batch_threads remains 3 (the same as the original attribute)

	  %1 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>
	  %2 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>
	  %3:2 = "tf.TPUCopyWithDynamicShape"(%1, %2, %cst, %cst) {operandSegmentSizes = array<i32: 2, 2>} : (tensor<2048xi32>, tensor<2048xi32>, tensor<i32>, tensor<i32>) -> (tensor<2048xi32>, tensor<2048xi32>)
	  tf_device.return %3#0, %3#1 : tensor<2048xi32>, tensor<2048xi32>

  // CHECK-NOT: tf.SendTPUEmbeddingGradients
  // CHECK-NOT: tf.RecvTPUEmbeddingActivations

  "tf.SendTPUEmbeddingGradients"(%arg0) {config = "\0A%\0A\0Dwatches_table\10\F5\03\18\80\02 \01*\0C\1A\00j\05\0D\00\00\80?\88\01\01\10\02\18\80\04 \01(\02", operandSegmentSizes = array<i32: 1, 0>} : (tensor<512x256xf32>) -> ()

    %0 = "tf.VarHandleOp"() {shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xi32>>>
    %1 = "tf.BatchFunction"(%0) {
        f = @called_by_batch_func,
        operandSegmentSizes = array<i32: 1, 0>,
        batch_timeout_micros = 1000,
        max_batch_size = 8,
        num_batch_threads = 2
      } : (tensor<!tf_type.resource<tensor<2xi32>>>) -> tensor<*xi32>
    return %1 : tensor<*xi32>

      %outputs_2, %control_2 = tf_executor.island wraps "tf.BatchFunction"(%outputs_1) {batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", f = @batched_func, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 0, 1>, shared_name = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
      tf_executor.fetch %outputs_2 : tensor<*xf32>
    }
    func.return %0 : tensor<*xf32>
  }

  %2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
  func.return %2 : tensor<1x3xf32>