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

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])