Convolution1DFlipout and DenseFlipout: apply_gradients() raise "TypeError: Can not convert a NoneType into a Tensor or Operation.".

[aikho]

When using tfp.layers.Convolution1DFlipout and tfp.layers.DenseFlipout in the same model calling apply_gradients() (using custom training loop and GradientTape()) raise error: TypeError: Can not convert a NoneType into a Tensor or Operation.

If there is just Convolution1DFlipout, it works. If there is just tfp.layers.DenseFlipout, it works.

tensorflow-gpu 2.0.0 tensorflow-probability 0.8.0

Minimal script to reproduce error:

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras.datasets import imdb
from tensorflow.keras.preprocessing import sequence

max_features = 5000
maxlen = 500
batch_size = 32
embedding_dims = 50
filters = 250
kernel_size = 3

(x_train, y_train), (x_val, y_val) = imdb.load_data(num_words=max_features)
x_train = sequence.pad_sequences(x_train, maxlen=maxlen)
x_val = sequence.pad_sequences(x_val, maxlen=maxlen)
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)
val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(batch_size)

model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Embedding(max_features,
                    embedding_dims,
                    input_length=maxlen))
model.add(tfp.layers.Convolution1DFlipout(filters,
                 kernel_size,
                 padding='valid',
                 activation='relu',
                 strides=1))
model.add(tf.keras.layers.GlobalMaxPooling1D())
model.add(tfp.layers.DenseFlipout(1, activation='sigmoid'))

optimizer = tf.keras.optimizers.Adam(lr=0.01)

for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
    with tf.GradientTape() as tape:
        logits = model(x_batch_train)
        labels_distribution = tfp.distributions.Bernoulli(logits=logits)
        loss_value = -tf.reduce_mean(labels_distribution.log_prob(y_batch_train))

        kl_loss = sum(model.losses) / float(len(x_train))
        loss_value = loss_value + kl_loss

    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))

Error trace:

TypeError                                 Traceback (most recent call last)
<ipython-input-2-2ef521648225> in <module>
     41 
     42     grads = tape.gradient(loss_value, model.trainable_weights)
---> 43     optimizer.apply_gradients(zip(grads, model.trainable_weights))

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py in apply_gradients(self, grads_and_vars, name)
    439           functools.partial(self._distributed_apply, apply_state=apply_state),
    440           args=(grads_and_vars,),
--> 441           kwargs={"name": name})
    442 
    443   def _distributed_apply(self, distribution, grads_and_vars, name, apply_state):

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/distribute/distribute_lib.py in merge_call(self, merge_fn, args, kwargs)
   1915     if kwargs is None:
   1916       kwargs = {}
-> 1917     return self._merge_call(merge_fn, args, kwargs)
   1918 
   1919   def _merge_call(self, merge_fn, args, kwargs):

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/distribute/distribute_lib.py in _merge_call(self, merge_fn, args, kwargs)
   1922         distribution_strategy_context._CrossReplicaThreadMode(self._strategy))  # pylint: disable=protected-access
   1923     try:
-> 1924       return merge_fn(self._strategy, *args, **kwargs)
   1925     finally:
   1926       _pop_per_thread_mode()

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py in _distributed_apply(self, distribution, grads_and_vars, name, apply_state)
    492         # context. (eager updates execute immediately)
    493         with ops._get_graph_from_inputs(update_ops).as_default():  # pylint: disable=protected-access
--> 494           with ops.control_dependencies(update_ops):
    495             return self._iterations.assign_add(1).op
    496 

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py in control_dependencies(control_inputs)
   5255     return NullContextmanager()
   5256   else:
-> 5257     return get_default_graph().control_dependencies(control_inputs)
   5258 
   5259 

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py in control_dependencies(self, control_inputs)
    354       else:
    355         filtered_control_inputs.append(graph_element)
--> 356     return super(FuncGraph, self).control_dependencies(filtered_control_inputs)
    357 
    358   def as_default(self):

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py in control_dependencies(self, control_inputs)
   4689           (hasattr(c, "_handle") and hasattr(c, "op"))):
   4690         c = c.op
-> 4691       c = self.as_graph_element(c)
   4692       if isinstance(c, Tensor):
   4693         c = c.op

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py in as_graph_element(self, obj, allow_tensor, allow_operation)
   3608 
   3609     with self._lock:
-> 3610       return self._as_graph_element_locked(obj, allow_tensor, allow_operation)
   3611 
   3612   def _as_graph_element_locked(self, obj, allow_tensor, allow_operation):

~/miniconda3/envs/nlp_tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py in _as_graph_element_locked(self, obj, allow_tensor, allow_operation)
   3697       # We give up!
   3698       raise TypeError("Can not convert a %s into a %s." %
-> 3699                       (type(obj).__name__, types_str))
   3700 
   3701   def get_operations(self):

TypeError: Can not convert a NoneType into a Tensor or Operation.

[RSamavedam]

I am having a similar error, were you able to solve this?