def METHOD_NAME(self):
    images = tf.ones((1, 2, 2, 1))
    sn_wrapper = keras.layers.SpectralNormalization(
        keras.layers.Conv2D(
            1, [2, 2], kernel_initializer=tf.constant_initializer(value=1)
        ),
        input_shape=(2, 2, 1),
    )
    result = sn_wrapper(images, training=False)
    result_train = sn_wrapper(images, training=True)
    expected_output = tf.constant([[[[4.0]]]], dtype=tf.float32)
    self.assertAllClose(result, expected_output)
    # max eigen value of 2x2 matrix of ones is 2
    self.assertAllClose(result_train, expected_output / 2)
    self.assertTrue(hasattr(sn_wrapper, "u"))