TypeError when using tf.cond()

General Discussion
, ,
#1

Below is the snippet I have used in Tensorflow 2.6.0

import tensorflow as tf
x = tf.keras.Input(shape=(224, 224, 3), batch_size=None)
x1=tf.keras.Input(1, dtype=tf.int32)
y = tf.cond(tf.less(x1,5), lambda :tf.keras.layers.ReLU()(x), lambda :tf.keras.layers.LeakyReLU(alpha=0.1)(x))
model=tf.keras.models.Model(inputs=[x,x1], outputs=[y])
model.summary()

Error:

tensorflow2/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 969, in convert
(str(python_func), type(x)))
TypeError: To be compatible with tf.eager.defun, Python functions must return zero or more Tensors; in compilation of <function at 0x7f9dd0148200>, found return value of type <class ‘keras.engine.keras_tensor.KerasTensor’>, which is not a Tensor.

#3

You can wrap tf.cond in a custom layer as described below

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):

  def call(self, inputs):
    x, x1 = inputs[0], inputs[1]
    return tf.cond(tf.less(x1,5), lambda :tf.keras.layers.ReLU()(x), lambda :tf.keras.layers.LeakyReLU(alpha=0.1)(x))

x = tf.keras.Input(shape=(224, 224, 3), batch_size=None)
x1= tf.keras.Input(1, dtype=tf.int32)
y = CustomLayer()([x, x1])
model=tf.keras.Model(inputs=[x,x1], outputs=[y])

Thank you.