[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
Loading
Please sign in to comment