def test_soft_round_layer_inverse_inverse_soft_rounds(self): alpha = 5.0 layer = soft_round.SoftRound(alpha=alpha, inverse=True) x = tf.linspace(-5.0, 5.0, num=50) y = layer(x) self.assertAllClose(y, soft_round_ops.soft_round_inverse(x, alpha=alpha))
def test_soft_round_inverse_values_and_gradients_are_finite(self, alpha): x = tf.linspace(-.5, .5, 11) # covers exact integers and half-integers with tf.GradientTape() as tape: tape.watch(x) y = soft_round_ops.soft_round_inverse(x, alpha=alpha) dy = tape.gradient(y, x) self.assertAllEqual(tf.math.is_finite(y), tf.ones(x.shape, dtype=bool)) if alpha > 15: # We allow non-finite values for large alphas, since the function simply # is extremely steep there. expected_finite = tf.one_hot(5, 11, False, True) else: expected_finite = tf.ones(x.shape, dtype=bool) self.assertAllEqual(tf.math.is_finite(dy), expected_finite)
def test_soft_round_inverse_values_and_gradients_are_finite(self, alpha): x = tf.linspace(-.5, .5, 11) # covers exact integers and half-integers with tf.GradientTape() as tape: tape.watch(x) y = soft_round_ops.soft_round_inverse(x, alpha=alpha) dy = tape.gradient(y, x) self.assertAllEqual(tf.math.is_finite(y), tf.ones(x.shape, dtype=bool)) is_finite = tf.math.is_finite(dy) expected_finite = tf.ones(dy.shape, dtype=bool) if alpha > 15: # We allow non-finite values at 0 for large alphas, since the function # simply is extremely steep there. expected_finite = tf.tensor_scatter_nd_update( expected_finite, [[5]], [is_finite[5]]) self.assertAllEqual(is_finite, expected_finite)
def inverse_transform(self, y): return soft_round_ops.soft_round_inverse(y, self._alpha)
def test_soft_round_inverse_large_alpha_is_ceil_minus_half(self): # We don't care what happens exactly near integer values: for offset in range(-5, 5): x = tf.linspace(offset + 0.001, offset + 0.999, 100) y = soft_round_ops.soft_round_inverse(x, alpha=5000.0) self.assertAllClose(tf.math.ceil(x) - 0.5, y, atol=0.001)
def test_soft_inverse_is_actual_inverse(self): x = tf.constant([-1.25, -0.75, 0.75, 1.25], dtype=tf.float32) y = soft_round_ops.soft_round(x, alpha=2.0) x2 = soft_round_ops.soft_round_inverse(y, alpha=2.0) self.assertAllClose(x, x2)
def test_soft_inverse_round_small_alpha_is_identity(self): x = tf.linspace(-2., 2., 50) y = soft_round_ops.soft_round_inverse(x, alpha=1e-13) self.assertAllEqual(x, y)