TensorFlow: Visualizing Custom Layers

Question

I\'m implementing a paper where I have to write custom layer. Following is the code.

class CouplingLayer(tf.keras.layers.Layer):
    """
    Ad