What is the TensorFlow v2 version of tf.compat.v1.Print?

Question

What is the TensorFlow v2 version of tf.compat.v1.Print? Specifically, is there a TF2 op that acts as the identity on its input, and prints a message when its output tensor is evaluated?

---

For context, I am using the Keras Functional API to define a model, similarly to:

input = tk.Input(input_shape)
layer_1 = tk.Dense(16)(input)
layer_2 = tk.Dense(16)(layer_1)
alpha = tf.math.exp(layer_2)
beta = alpha**2
model = tk.Model(inputs=[input], outputs=[beta])

I would like to examine the intermediate values of the model (for example, alpha) at inference time. I can do this as follows:

input = tk.Input(input_shape)
layer_1 = tk.Dense(16)(input)
layer_2 = tk.Dense(16)(layer_1)
alpha = tf.math.exp(layer_2)
alpha = tf.compat.v1.Print(alpha, [f'alpha={alpha}'])
beta = alpha**2
model = tk.Model(inputs=[input], outputs=[beta])

Is there a TF2 way of doing this?

(I know I could return alpha as an output, but this would require touching every place that an instance of this model is called, which is annoying.)

Answer 1

Based on this answer, creating a Custom Layer with the print functionality will do the trick.

A minimal example:

class PrintLayer(tf.keras.layers.Layer):
    def __init__(self, trainable=False, **kwargs):
        super().__init__(trainable=trainable, **kwargs)

    def call(self, inputs):
        tf.print(inputs, output_stream=sys.stdout)
        return inputs

input = tk.Input(input_shape)
layer_1 = tk.Dense(16)(input)
layer_2 = tk.Dense(16)(layer_1)
alpha = tf.math.exp(layer_2)
alpha = PrintLayer()(alpha)
beta = alpha**2
model = tk.Model(inputs=[input], outputs=[beta])