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)