A tensorflow layer that prints a warning mesage if there is a `nan` or `inf` in a layer's output.

check_nan.py

@tf.function
def has_nan(X):
    return tf.math.reduce_any(tf.math.is_nan(X))

@tf.function
def has_inf(X):
    return tf.math.reduce_any(tf.math.is_inf(X))

def check_nan(name):
    @tf.function
    def _inner(X):
        if has_nan(X) or has_inf(X):
            tf.print('nan/inf in', name,'->',X,'\n')

        return X
    
    return _inner

...
X = tf.keras.layers.Lambda(check_nan('layer name'))(X)
...