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)