def testGeneratorDatasetFinalizeFunctionCalled(self): # NOTE(mrry): This test tests the internal `_GeneratorDataset`, # which affords more control over what the finalize function can do than # the `Dataset.from_generator()` wrapper. # Use an `Event` to signal that the generator has been deleted. event = threading.Event() def finalize_fn(_): def finalize_py_func(): event.set() return 0 return script_ops.py_func(finalize_py_func, [], [dtypes.int64], stateful=True) dummy = constant_op.constant(37) dataset = dataset_ops._GeneratorDataset( dummy, lambda x: x, lambda x: x, finalize_fn, tensor_spec.TensorSpec((), dtypes.int32)) dataset = dataset.take(2) get_next = self.getNext(dataset) self.assertAllEqual(37, self.evaluate(get_next())) self.assertAllEqual(37, self.evaluate(get_next())) with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next())
def testGeneratorDatasetFinalizeFunctionCalled(self): # NOTE(mrry): This test tests the internal `_GeneratorDataset`, # which affords more control over what the finalize function can do than # the `Dataset.from_generator()` wrapper. # Use an `Event` to signal that the generator has been deleted. event = threading.Event() def finalize_fn(_): def finalize_py_func(): event.set() return 0 return script_ops.py_func(finalize_py_func, [], [dtypes.int64], stateful=True) dummy = constant_op.constant(37) iterator = (dataset_ops._GeneratorDataset(dummy, lambda x: x, lambda x: x, finalize_fn) .take(2) .make_initializable_iterator()) init_op = iterator.initializer get_next = iterator.get_next() with self.cached_session() as sess: sess.run(init_op) self.assertAllEqual(37, sess.run(get_next)) self.assertAllEqual(37, sess.run(get_next)) with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) self.assertTrue(event.is_set())
def testGeneratorDatasetFinalizeFunctionCalled(self): # NOTE(mrry): This test tests the internal `_GeneratorDataset`, # which affords more control over what the finalize function can do than # the `Dataset.from_generator()` wrapper. # Use an `Event` to signal that the generator has been deleted. event = threading.Event() def finalize_fn(_): def finalize_py_func(): event.set() return 0 return script_ops.py_func(finalize_py_func, [], [dtypes.int64], stateful=True) dummy = constant_op.constant(37) iterator = (dataset_ops._GeneratorDataset( dummy, lambda x: x, lambda x: x, finalize_fn).take(2).make_initializable_iterator()) init_op = iterator.initializer get_next = iterator.get_next() with self.test_session() as sess: sess.run(init_op) self.assertAllEqual(37, sess.run(get_next)) self.assertAllEqual(37, sess.run(get_next)) with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) self.assertTrue(event.is_set())
def fn(): output_signature = tensor_spec.TensorSpec((), dtypes.int64) dataset = dataset_ops._GeneratorDataset(1, init_fn, next_fn, finalize_fn, output_signature) iterator = iter(dataset) next(iterator)
def fn(): dataset = dataset_ops._GeneratorDataset( 1, init_fn, next_fn, finalize_fn, output_signature=tensor_spec.TensorSpec([], dtypes.int64)) iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator( dataset, [self._devices[1], self._devices[2]]) next(iterator)
def fn(): dataset = dataset_ops._GeneratorDataset(1, init_fn, next_fn, finalize_fn) iterator = iter(dataset) next(iterator)
def fn(): dataset = dataset_ops._GeneratorDataset(1, init_fn, next_fn, finalize_fn) iterator = multi_device_iterator_ops.MultiDeviceIteratorV2( dataset, ["/cpu:0", "/gpu:0"]) next(iterator)
def fn(): dataset = dataset_ops._GeneratorDataset(1, init_fn, next_fn, finalize_fn) iterator = iter(dataset) next(iterator)