def test_soft_round_values_and_gradients_are_finite(self, alpha):
x = tf.linspace(0., 1., 11) # covers exact integers and half-integers
with tf.GradientTape() as tape:
tape.watch(x)
y = soft_round_ops.soft_round(x, alpha=alpha)
dy = tape.gradient(y, x)
self.assertAllEqual(tf.math.is_finite(y), tf.ones(x.shape, dtype=bool))
self.assertAllEqual(tf.math.is_finite(dy), tf.ones(x.shape, dtype=bool))
class AdaptersTest(tf.test.TestCase, parameterized.TestCase):
@parameterized.named_parameters(
("softround_deepfactorized",
lambda d: round_adapters.SoftRoundAdapter(d, alpha=5.0),
deep_factorized.DeepFactorized, 0.0),
("softround_logistic",
lambda d: round_adapters.SoftRoundAdapter(d, alpha=5.0),
lambda: tfp.distributions.Logistic(loc=10.3, scale=1.5),
lambda: soft_round_ops.soft_round(0.3, alpha=5.0)),
("softround_normal",
lambda d: round_adapters.SoftRoundAdapter(d, alpha=4.0),
lambda: tfp.distributions.Normal(loc=10.4, scale=1.5),
lambda: soft_round_ops.soft_round(0.4, alpha=4.0)),
("noisysoftround_deepfactorized",
lambda d: round_adapters.NoisySoftRoundAdapter(d, alpha=5.0),
deep_factorized.DeepFactorized, 0.0),
("noisysoftround_logistic",
lambda d: round_adapters.NoisySoftRoundAdapter(d, alpha=5.0),
lambda: tfp.distributions.Logistic(loc=10, scale=1.5), 0.0),
("noisysoftround_normal",
lambda d: round_adapters.NoisySoftRoundAdapter(d, alpha=5.0),
lambda: tfp.distributions.Normal(loc=10, scale=1.5), 0.0),
("round_deepfactorized", round_adapters.RoundAdapter,
lambda: deep_factorized.DeepFactorized(init_scale=1.0), 0.0),
("round_logistic", round_adapters.RoundAdapter,
lambda: tfp.distributions.Logistic(loc=1.5, scale=1.5), 0.0),
("round_normal", round_adapters.RoundAdapter,
lambda: tfp.distributions.Normal(loc=1.5, scale=1.5), 0.0),
("noisyround_deepfactorized", round_adapters.NoisyRoundAdapter,
lambda: deep_factorized.DeepFactorized(init_scale=1.0), 0.0),
("noisyround_logistic", round_adapters.NoisyRoundAdapter,
lambda: tfp.distributions.Logistic(loc=1.5, scale=1.5), 0.0),
("noisyround_normal", round_adapters.NoisyRoundAdapter,
lambda: tfp.distributions.Normal(loc=1.5, scale=1.5), 0.0),
)
def test_tails_and_offset(self, adapter, distribution, expected_offset):
dist = adapter(distribution())
lower_tail = dist._lower_tail(2**-8)
try:
left_mass = dist.cdf(lower_tail)
except NotImplementedError:
# We use base distribution as a proxy for the tail mass.
left_mass = dist.base.cdf(lower_tail)
self.assertLessEqual(left_mass, 2**-8)
upper_tail = dist._upper_tail(2**-8)
try:
right_mass = dist.survival_function(upper_tail)
except NotImplementedError:
# We use base distribution as a proxy for the tail mass.
right_mass = dist.base.survival_function(upper_tail)
self.assertLessEqual(right_mass, 2**-8)
self.assertGreater(upper_tail, lower_tail)
offset = dist._quantization_offset()
if not isinstance(expected_offset, float):
# We cannot run tf inside the parameterized test declaration, hence
# non-float values are wrapped in a lambda.
expected_offset = expected_offset()
self.assertAllClose(offset, expected_offset)
@parameterized.named_parameters(
("softround_logistic",
lambda d: round_adapters.SoftRoundAdapter(d, alpha=5.0),
lambda: tfp.distributions.Logistic(loc=10, scale=1.5)),
("softround_normal",
lambda d: round_adapters.SoftRoundAdapter(d, alpha=5.0),
lambda: tfp.distributions.Normal(loc=10, scale=1.5)),
)
def test_mode_and_quantile(self, adapter, distribution):
dist = adapter(distribution())
mode = dist.mode()
left_mass = dist.cdf(mode)
self.assertAllClose(left_mass, 0.5)
quantile_75p = dist.quantile(0.75)
left_mass = dist.cdf(quantile_75p)
self.assertAllClose(left_mass, 0.75)
def test_lacking_mode_and_quantile(self):
dist = round_adapters.RoundAdapter(
tfp.distributions.Logistic(loc=1.5, scale=1.5))
with self.assertRaises(NotImplementedError):
dist.mode()
with self.assertRaises(NotImplementedError):
dist.quantile(0.75)
def test_lacking_tails_and_offset(self):
class NonInvertibleAdapater(round_adapters.MonotonicAdapter):
invertible = False
def transform(self, x):
return tf.ceil(x)
def inverse_transform(self, y):
return tf.floor(y)
dist = NonInvertibleAdapater(
tfp.distributions.Normal(loc=1.5, scale=1.5))
with self.assertRaises(NotImplementedError):
dist._lower_tail(0.01)
with self.assertRaises(NotImplementedError):
dist._upper_tail(0.01)
def transform(self, x):
return soft_round_ops.soft_round(x, self._alpha)
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_round_large_alpha_is_round(self):
# We don't care what happens exactly near half-integer values:
for offset in range(-5, 5):
x = tf.linspace(offset - 0.499, offset + 0.499, 100)
y = soft_round_ops.soft_round(x, alpha=2000.0)
self.assertAllClose(tf.round(x), y, atol=0.02)
def test_soft_round_small_alpha_is_identity(self): x = tf.linspace(-2., 2., 50) y = soft_round_ops.soft_round(x, alpha=1e-13) self.assertAllClose(x, y)
def test_soft_round_layer_soft_rounds(self):
alpha = 5.0
layer = soft_round.SoftRound(alpha=alpha)
x = tf.linspace(-5.0, 5.0, num=50)
y = layer(x)
self.assertAllClose(y, soft_round_ops.soft_round(x, alpha=alpha))
