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(...).