diff options
Diffstat (limited to 'tensorflow/python/data/experimental/ops/random_ops.py')
| -rw-r--r-- | tensorflow/python/data/experimental/ops/random_ops.py | 21 |
1 files changed, 17 insertions, 4 deletions
diff --git a/tensorflow/python/data/experimental/ops/random_ops.py b/tensorflow/python/data/experimental/ops/random_ops.py index e3a2aeab31..25d7fbf691 100644 --- a/tensorflow/python/data/experimental/ops/random_ops.py +++ b/tensorflow/python/data/experimental/ops/random_ops.py @@ -33,13 +33,26 @@ class RandomDataset(dataset_ops.DatasetSource): def __init__(self, seed=None): """A `Dataset` of pseudorandom values.""" super(RandomDataset, self).__init__() - self._seed, self._seed2 = random_seed.get_seed(seed) + + # NOTE(mrry): We generate the seed-pair once per graph in which the dataset + # is iterated over, and cache it in `self._graph_seed_map`. This supports + # two features: iterating over the same `ShuffleDataset` twice in the same + # pipeline and observing the same order (by tying the seeds together with + # a randomly-generated seed), and using `Dataset.make_one_shot_iterator()`, + # which requires the stateful RNG op to be created inside the same graph as + # the dataset. + self._original_seed = seed + self._graph_seed_map = {} def _as_variant_tensor(self): + try: + seed, seed2 = self._graph_seed_map[ops.get_default_graph()] + except KeyError: + seed, seed2 = random_seed.get_seed(self._original_seed) + self._graph_seed_map[ops.get_default_graph()] = (seed, seed2) + return gen_dataset_ops.random_dataset( - seed=self._seed, - seed2=self._seed2, - **dataset_ops.flat_structure(self)) + seed=seed, seed2=seed2, **dataset_ops.flat_structure(self)) @property def output_classes(self): |