Commit be2b6e19 authored by Derek Murray's avatar Derek Murray Committed by TensorFlower Gardener
Browse files

[tf.data] Add optional `args` argument to `Dataset.from_generator()`.

The new argument allows you to parameterize the generator with the value of a tf.Tensor,
enabling `Dataset.from_generator()` to be initialized from a placeholder or used in a
nested expression (such as `flat_map()` or `parallel_interleave()`). For example:

```python
def generator(n):
  for _ in range(n):
    yield n

# Define a generator based on a placeholder.
placeholder = tf.placeholder(tf.int64, shape=[])
dataset = tf.data.Dataset.from_generator(generator, tf.int64, args=(placeholder,))

# Define a generator based on the value of a nested dataset element.
dataset = tf.data.Dataset.range(10).flat_map(
    lambda i: tf.data.Dataset.from_generator(generator, tf.int64, args=(i,)))
```

Fixes #19269. Partially addresses issue #13101.

PiperOrigin-RevId: 196598650
parent 4c2cb712
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment