I'm trying to get a preprocessing function to work with the Dataset map, but I get the following error (full stack trace at the bottom):

```
ValueError: Tensor-typed variable initializers must either be wrapped in an init_scope or callable (e.g., `tf.Variable(lambda : tf.truncated_normal([10, 40]))`) when building functions. Please file a feature request if this restriction inconveniences you.

```

Below is a full snippet that reproduces the issue. My question is, why in one use case (crop only) it works, and when RandomFlip is used it doesn't? How can this be fixed?

```py
import functools
import numpy as np
import tensorflow as tf

def data_gen():
    for i in range(10):
        x = np.random.random(size=(80, 80, 3)) * 255  # rgb image
        x = x.astype('uint8')
        y = np.random.random(size=(40, 40, 1)) * 255  # downsized mono image
        y = y.astype('uint8')
        yield x, y

def preprocess(image, label, cropped_image_size, cropped_label_size, skip_augmentations=False):

    x = image
    y = label

    x_size = cropped_image_size
    y_size = cropped_label_size

    if not skip_augmentations:
        x = tf.keras.layers.RandomFlip(mode="horizontal")(x)
        y = tf.keras.layers.RandomFlip(mode="horizontal")(y)

        x = tf.keras.layers.RandomRotation(factor=1.0, fill_mode='constant')(x)
        y = tf.keras.layers.RandomRotation(factor=1.0, fill_mode='constant')(y)

    x = tf.keras.layers.CenterCrop(x_size, x_size)(x)
    y = tf.keras.layers.CenterCrop(y_size, y_size)(y)

    return x, y

print(tf.__version__) # 2.6.0
dataset = tf.data.Dataset.from_generator(data_gen, output_signature=(
    tf.TensorSpec(shape=(80, 80, 3), dtype='uint8'),
    tf.TensorSpec(shape=(40, 40, 1), dtype='uint8')
))

crop_only_fn = functools.partial(preprocess, cropped_image_size=50, cropped_label_size=25, skip_augmentations=True)
train_preprocess_fn = functools.partial(preprocess, cropped_image_size=50, cropped_label_size=25, skip_augmentations=False)

# This works
crop_dataset = dataset.map(crop_only_fn)

# This fails: ValueError: Tensor-typed variable initializers must either be wrapped in an init_scope or callable
train_dataset = dataset.map(train_preprocess_fn)

```

Full-stack trace:

```
Traceback (most recent call last):
  File "./issue_dataaug.py", line 50, in <module>
    train_dataset = dataset.map(train_preprocess_fn)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 1861, in map
    return MapDataset(self, map_func, preserve_cardinality=True)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 4985, in __init__
    use_legacy_function=use_legacy_function)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 4218, in __init__
    self._function = fn_factory()
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 3151, in get_concrete_function
    *args, **kwargs)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 3116, in _get_concrete_function_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 3463, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 3308, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 1007, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 4195, in wrapped_fn
    ret = wrapper_helper(*args)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 4125, in wrapper_helper
    ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
  File "/...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 695, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    ./issue_dataaug.py:25 preprocess  *
        x = tf.keras.layers.RandomFlip(mode="horizontal")(x)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/keras/layers/preprocessing/image_preprocessing.py:414 __init__  **
        self._rng = make_generator(self.seed)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/keras/layers/preprocessing/image_preprocessing.py:1375 make_generator
        return tf.random.Generator.from_non_deterministic_state()
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/stateful_random_ops.py:396 from_non_deterministic_state
        return cls(state=state, alg=alg)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/stateful_random_ops.py:476 __init__
        trainable=False)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/stateful_random_ops.py:489 _create_variable
        return variables.Variable(*args, **kwargs)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:268 __call__
        return cls._variable_v2_call(*args, **kwargs)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:262 _variable_v2_call
        shape=shape)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:243 <lambda>
        previous_getter = lambda **kws: default_variable_creator_v2(None, **kws)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/variable_scope.py:2675 default_variable_creator_v2
        shape=shape)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:270 __call__
        return super(VariableMetaclass, cls).__call__(*args, **kwargs)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:1613 __init__
        distribute_strategy=distribute_strategy)
    /...//virtualenvs/cvi36/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:1695 _init_from_args
        raise ValueError("Tensor-typed variable initializers must either be "

    ValueError: Tensor-typed variable initializers must either be wrapped in an init_scope or callable (e.g., `tf.Variable(lambda : tf.truncated_normal([10, 40]))`) when building functions. Please file a feature request if this restriction inconveniences you.

```