}) {n = 2 : i32, operandSegmentSizes = array<i32: 0, 0>} : () -> ()
  func.return
}

// -----

// Check that an empty replicate block is invalid.
func.func @verifier_replicate_empty_block() {
  "tf_device.replicate" () ({
// expected-error@-1 {{'tf_device.replicate' op expects a non-empty block}}
  ^entry:
  }) {n = 2 : i32, operandSegmentSizes = array<i32: 0, 0>} : () -> ()
  func.return
}

  %1 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
  func.return %1 : tensor<i32>
}

    // 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>) -> ()

  // CHECK: "tf.BatchFunction"(%arg0, %0)
  // CHECK: operandSegmentSizes = array<i32: 1, 1>