I have some code written with tensorflow 2.3 in which tensors are constructed using tf.keras. As moving to tensorflow > 2.16, tf.keras is retired an tf calls keras >= 3.0. The type of the tensors in my code changes from tf.Tensor to Keras.Tensor. As a consequence, my code is largely affected, most of which raises the following error ValueError: A KerasTensor cannot be used as input to a TensorFlow function. A KerasTensor is a symbolic placeholder for a shape and dtype, used when constructing Keras Functional models or Keras Functions. You can only use it as input to a Keras layer or a Keras operation (from the namespaces keras.layers and keras.operations). You are likely doing something like: x = Input(...) ... tf_fn(x) # Invalid. What you should do instead is wrap tf_fn in a layer: class MyLayer(Layer): def call(self, x): return tf_fn(x) x = MyLayer()(x) So is there a simple solution to solve this or do I need to apply this wrapper everywhere? Update 2024-09-11 Let me try to make the question more solid. Consider the following example. It is taken from function get_decoder_mask in google tft module In TF 2.3, using tf.keras, we can define the following tenor for masking as follows. Define t = keras.layers.Input(shape=(250, 5), name="input") which is tf.Tensor. len_s = tf.shape(t)[-2] bs = tf.shape(t)[:-2] mask = tf.cumsum(tf.eye(len_s, batch_shape=bs), -2) In TF 2.17, only native keras is available. t = keras.layers.Input(shape=(250, 5), name="input") which is KerasTensor. Then the following code does not work. len_s = tf.shape(t)[-2] bs = tf.shape(t)[:-2] mask = tf.cumsum(tf.eye(len_s, batch_shape=bs), -2) With tf.shape(t) being replaced with t.shape, the code breaks further down at tf.eye(...). Continue reading...