tensorflow: tf.summary.image does not work with keras layers in tf 2.0 due to EagerExecution issues

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow):Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Official tf docker
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary):No
• TensorFlow version (use command below):2.0.0-alpha0
• Python version: 3.5
• CUDA/cuDNN version: 10.1
• GPU model and memory:Titan V

I use tf.keras subclassing, but when I want to track the some images in my network, the error at the bottom of the post occurs. Initially, I think the reason for the error is, that tf.summary.image becomes a normal Tensor as input, but needs an eagerTensor. If I create a random tensor in the summary function and evaluate the type it is an eagerTensor, but if I use a Tensor from a keras.layers.Layer class it evaluates to a normal Tensor and throws the error (you can try it in your code).

First I thought that the tf.dataset returns a non-eager tensor, but if I use a numpy array as input, the same error occurs as well. So the error might be caused by the keras.layer.Layer class.

The error occurs on GPU and CPU.

Working example

import tensorflow as tf
import numpy as np

inp_arr_without_ds = np.random.rand(2, 200, 200, 3)

inp_arr_4_ds = np.random.rand(2, 2, 200, 200, 3)
tf_ds = tf.data.Dataset.from_tensor_slices((inp_arr_4_ds, inp_arr_4_ds))
tf_ds = tf_ds.map(lambda x, y: (x, y)).repeat(10).shuffle(10)

class random_model(tf.keras.Model):
    
    def __init__(self, name):
        super(random_model, self).__init__(name=name)
        self.conv_1 = tf.keras.layers.Conv2D(3, [3, 3], padding="same")
        self.tf_board_writer = tf.summary.create_file_writer("test")
        self.img_callback = [tf.keras.callbacks.LambdaCallback(on_epoch_end=self.save_img)]
        
    def call(self, inputs):
        self.initialized_layer = self.conv_1(inputs)
        return self.initialized_layer
    
    def save_img(self, epochs, logs):
        with self.tf_board_writer.as_default():
            # Does not work: type: class 'tensorflow.python.framework.ops.Tensor' 
            tf.summary.image("image", self.initialized_layer, step=epochs)
            
            # Does work type: class 'tensorflow.python.framework.ops.EagerTensor'
            #tf.summary.image("image", tf.random.uniform([2, 200, 200, 3]), step=epochs) 
            
    def compile_model(self):
        self.compile(tf.optimizers.Adam(0.001), tf.losses.mean_absolute_error)
    
    def fit_model_with_ds(self, ds):
        self.fit(ds, callbacks=self.img_callback)
        
    def fit_model_with_array(self, x, y):
        self.fit(x, y, callbacks=self.img_callback)

print(tf.__version__)        

# Both do not work
non_ds_model = random_model("non_ds") 
non_ds_model.compile_model()
non_ds_model.fit_model_with_array(inp_arr_without_ds, inp_arr_without_ds)
            
tf_ds_model = random_model("tf_ds")
tf_ds_model.compile_model()
tf_ds_model.fit_model_with_ds(tf_ds)

Error Message

> 2.0.0-alpha0
> 
> ---------------------------------------------------------------------------
> TypeError                                 Traceback (most recent call last)
> <ipython-input-25-0b5a4b30a43a> in <module>
>      39 non_ds_model = random_model("non_ds")
>      40 non_ds_model.compile_model()
> ---> 41 non_ds_model.fit_model_with_array(inp_arr_without_ds, inp_arr_without_ds)
>      42 
>      43 tf_ds_model = random_model("tf_ds")
> 
> <ipython-input-25-0b5a4b30a43a> in fit_model_with_array(self, x, y)
>      32 
>      33     def fit_model_with_array(self, x, y):
> ---> 34         self.fit(x, y, callbacks=self.img_callback)
>      35 
>      36 print(tf.__version__)
> 
> /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, max_queue_size, workers, use_multiprocessing, **kwargs)
>     871           validation_steps=validation_steps,
>     872           validation_freq=validation_freq,
> --> 873           steps_name='steps_per_epoch')
>     874 
>     875   def evaluate(self,
> 
> /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training_arrays.py in model_iteration(model, inputs, targets, sample_weights, batch_size, epochs, verbose, callbacks, val_inputs, val_targets, val_sample_weights, shuffle, initial_epoch, steps_per_epoch, validation_steps, validation_freq, mode, validation_in_fit, prepared_feed_values_from_dataset, steps_name, **kwargs)
>     406     if mode == ModeKeys.TRAIN:
>     407       # Epochs only apply to `fit`.
> --> 408       callbacks.on_epoch_end(epoch, epoch_logs)
>     409     progbar.on_epoch_end(epoch, epoch_logs)
>     410 
> 
> /usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/callbacks.py in on_epoch_end(self, epoch, logs)
>     288     logs = logs or {}
>     289     for callback in self.callbacks:
> --> 290       callback.on_epoch_end(epoch, logs)
>     291 
>     292   def on_train_batch_begin(self, batch, logs=None):
> 
> <ipython-input-25-0b5a4b30a43a> in save_img(self, epochs, logs)
>      20         with self.tf_board_writer.as_default():
>      21             # Does not work: type: class 'tensorflow.python.framework.ops.Tensor'
> ---> 22             tf.summary.image("image", self.initialized_layer, step=epochs)
>      23 
>      24             # Does work type: class 'tensorflow.python.framework.ops.EagerTensor'
> 
> /usr/local/lib/python3.5/dist-packages/tensorboard/plugins/image/summary_v2.py in image(name, data, step, max_outputs, description)
>      71     encoded_images = tf.map_fn(tf.image.encode_png, limited_images,
>      72                                dtype=tf.string,
> ---> 73                                name='encode_each_image')
>      74     # Workaround for map_fn returning float dtype for an empty elems input.
>      75     encoded_images = tf.cond(
> 
> /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/map_fn.py in map_fn(fn, elems, dtype, parallel_iterations, back_prop, swap_memory, infer_shape, name)
>     226                                      dynamic_size=False,
>     227                                      infer_shape=True)
> --> 228         for elem in elems_flat]
>     229     # Unpack elements
>     230     elems_ta = [
> 
> /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/map_fn.py in <listcomp>(.0)
>     226                                      dynamic_size=False,
>     227                                      infer_shape=True)
> --> 228         for elem in elems_flat]
>     229     # Unpack elements
>     230     elems_ta = [
> 
> /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/tensor_array_ops.py in __init__(self, dtype, size, dynamic_size, clear_after_read, tensor_array_name, handle, flow, infer_shape, element_shape, colocate_with_first_write_call, name)
>    1036         element_shape=element_shape,
>    1037         colocate_with_first_write_call=colocate_with_first_write_call,
> -> 1038         name=name)
>    1039 
>    1040     self._implementation.parent = weakref.ref(self)
> 
> /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/tensor_array_ops.py in __init__(***failed resolving arguments***)
>     742     if isinstance(size, ops.EagerTensor):
>     743       size = size.numpy()
> --> 744     self._tensor_array = [None for _ in range(size)]
>     745 
>     746   @property
> 
> TypeError: 'Tensor' object cannot be interpreted as an integer
>