def testHardenAttentionAllZeros(self, gumbel_noise_weight):
    """Check if the hardening code does not divide by zero for all zeros."""
    x = np.zeros((5, 3, 12), dtype=np.float32)
    y = common_attention.harden_attention_weights(
        tf.constant(x, dtype=tf.float32), 3, gumbel_noise_weight)
    res = self.evaluate(y)
    if gumbel_noise_weight <= 0.0:
      self.assertAllClose(res, x)
Check if the hardening code does not divide by zero for all zeros.
tensor2tensor/layers/common_attention_test.py