class_docstr: Tests for tf.data placement within tf.functions.
  
  Specifically, tf.data dataset tensors cannot be copied between devices. These
  tests verify the ops are placed in a way that avoids this.
class_docstr_tok: ["Tests","for","tf.data","placement","within","tf.functions",".","Specifically",",","tf.data","dataset","tensors","can","not","be","copied","between","devices",".","These","tests","verify","the","ops","are","placed","in","a","way","that","avoids","this","."]
code: def testCondWithColocation(self):
      # When the cond op is colocated with the dataset, there shouldn't be
      # cross-device copies.
      @def_function.function
      def f():
          dataset = dataset_ops.Dataset.range(8)
  
          def fn():
              return dataset.map(lambda x: x + 1)
  
          c = constant_op.constant(2)
          with ops.colocate_with(
              dataset._variant_tensor
          ):  # pylint:disable=protected-access
              a = cond.cond(math_ops.equal(c, 2), fn, fn)
              iterator = iter(a)
              nxt = next(iterator)
          return nxt
  
      self.assertEqual(f().numpy(), 1)
  
imports: from absl.testing import parameterized
  from tensorflow.python.data.experimental.ops import prefetching_ops
  from tensorflow.python.data.kernel_tests import test_base
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import cond
  from tensorflow.python.ops import gen_dataset_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 5 for calling the PlacementTest obj's underlying member method code with example usage: obj.testCondWithColocation() and returns: nxt, dataset