def METHOD_NAME(x):
    with tf.GradientTape() as tape:
        y = layer(x)
        loss = tf.math.reduce_sum(y)
    gradients = tape.gradient(loss, layer.weights)
    self.assertListEqual(gradients[0].shape.as_list(), [50, 16])