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, 4 insertions, 17 deletions
diff --git a/tensorflow/python/data/experimental/ops/random_ops.py b/tensorflow/python/data/experimental/ops/random_ops.py index 25d7fbf691..e3a2aeab31 100644 --- a/tensorflow/python/data/experimental/ops/random_ops.py +++ b/tensorflow/python/data/experimental/ops/random_ops.py @@ -33,26 +33,13 @@ class RandomDataset(dataset_ops.DatasetSource): def __init__(self, seed=None): """A `Dataset` of pseudorandom values.""" super(RandomDataset, self).__init__() - - # 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 = {} + self._seed, self._seed2 = random_seed.get_seed(seed) 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=seed, seed2=seed2, **dataset_ops.flat_structure(self)) + seed=self._seed, + seed2=self._seed2, + **dataset_ops.flat_structure(self)) @property def output_classes(self): |