def testJacobian(self):
bijector = tfb.Exp()
x = tf.constant([22.])
fldj = bijector.forward_log_det_jacobian(x, event_ndims=0)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=0)
fldj_, fldj_theoretical_ = self.evaluate([fldj, fldj_theoretical])
self.assertAllClose(fldj_, fldj_theoretical_)
x = np.expand_dims(np.linspace(-1, 1, num=10), -1)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=1)
fldj_, fldj_theoretical_ = self.evaluate([fldj, fldj_theoretical])
self.assertAllClose(fldj_, fldj_theoretical_)
def testTheoreticalFldj(self):
raw_mat = tf.constant([[1., 2, 3],
[4, 5, 6],
[0.5, 0., 0.25]])
nbatch = 5
batch_mats = raw_mat * tf.range(1., nbatch + 1.)[:, tf.newaxis, tf.newaxis]
lower_upper, permutation = tf.linalg.lu(tf.cast(batch_mats, tf.float64))
bijector = tfb.ScaleMatvecLU(
lower_upper=lower_upper, permutation=permutation, validate_args=True)
self.assertEqual(tf.float64, bijector.dtype)
channels = tf.compat.dimension_value(lower_upper.shape[-1])
x = np.random.uniform(size=[2, 7, nbatch, channels]).astype(np.float64)
y = self.evaluate(bijector.forward(x))
bijector_test_util.assert_bijective_and_finite(
bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=1,
inverse_event_ndims=1,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
# The jacobian is not yet broadcast, since it is constant.
fldj = fldj + tf.zeros(tf.shape(x)[:-1], dtype=x.dtype)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=1)
self.assertAllClose(
self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldjSimple(self):
bijector = tfb.RationalQuadraticSpline(
bin_widths=[1., 1],
bin_heights=[np.sqrt(.5), 2 - np.sqrt(.5)],
knot_slopes=1)
self.assertEqual(tf.float64, bijector.dtype)
dim = 5
x = np.linspace(-1.05, 1.05, num=2 * dim, dtype=np.float64).reshape(2, dim)
y = self.evaluate(bijector.forward(x))
bijector_test_util.assert_bijective_and_finite(
bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=0,
inverse_event_ndims=0,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=0)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=0)
self.assertAllClose(
self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldj(self, data):
# get_fldj_theoretical test rig requires 1-d batches.
batch_shape = data.draw(tfp_hps.shapes(min_ndims=1, max_ndims=1))
bijector = data.draw(
rq_splines(batch_shape=batch_shape, dtype=tf.float64))
self.assertEqual(tf.float64, bijector.dtype)
bw, bh, kd = self.evaluate(
[bijector.bin_widths, bijector.bin_heights, bijector.knot_slopes])
logging.info('bw: %s\nbh: %s\nkd: %s', bw, bh, kd)
x_shp = ((bw + bh)[..., :-1] + kd).shape[:-1]
if x_shp[-1] == 1: # Possibly broadcast the x dim.
dim = data.draw(hps.integers(min_value=1, max_value=7))
x_shp = x_shp[:-1] + (dim, )
x = np.linspace(-5, 5, np.prod(x_shp),
dtype=np.float64).reshape(*x_shp)
y = self.evaluate(bijector.forward(x))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=0,
inverse_event_ndims=0,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=0)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=0)
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testWithLKJSamples(self, dimension, concentration):
bijector = tfb.CorrelationCholesky()
lkj_dist = lkj.LKJ(dimension=dimension,
concentration=np.float64(concentration),
input_output_cholesky=True)
batch_size = 10
y = self.evaluate(lkj_dist.sample([batch_size]))
x = self.evaluate(bijector.inverse(y))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=1,
inverse_event_ndims=2,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector,
x,
event_ndims=1,
inverse_event_ndims=2,
output_to_unconstrained=tfb.Invert(tfb.FillTriangular()))
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldj(self, data):
dim = data.draw(hps.integers(min_value=0, max_value=10))
diag_bijector = data.draw(
bijector_hps.unconstrained_bijectors(
max_forward_event_ndims=1,
must_preserve_event_ndims=True).filter(
_preserves_vector_dim(dim)))
logging.info('Using diagonal bijector %s %s', diag_bijector.name,
diag_bijector)
bijector = tfb.TransformDiagonal(diag_bijector=diag_bijector)
ensure_nonzero_batch = lambda shape: [d if d > 0 else 1 for d in shape]
shape = data.draw(
tfp_hps.shapes().map(ensure_nonzero_batch)) + [dim, dim]
x = np.random.randn(*shape).astype(np.float64)
y = self.evaluate(bijector.forward(x))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=2,
inverse_event_ndims=2,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=2)
# For constant-jacobian bijectors, the zero fldj may not be broadcast.
fldj = fldj + tf.zeros(tf.shape(x)[:-2], dtype=x.dtype)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=2, inverse_event_ndims=2)
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldj(self):
width = 4
bijector = tfp.experimental.bijectors.build_trainable_highway_flow(
width,
activation_fn=tf.nn.softplus,
gate_first_n=2,
seed=test_util.test_seed())
self.evaluate([v.initializer for v in bijector.trainable_variables])
x = self.evaluate(
samplers.uniform(
[width],
minval=-1.,
maxval=1.,
seed=test_util.test_seed(sampler_type='stateless')))
y = self.evaluate(bijector.forward(x))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=1,
inverse_event_ndims=1,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
# The jacobian is not yet broadcast, since it is constant.
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=1)
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldj(self):
nbatch = 5
channels = 10
x = np.random.uniform(size=[nbatch, channels]).astype(np.float32)
bijector = tfb.BatchNormalization(training=False)
bijector.batchnorm.build(x.shape)
self.evaluate([v.initializer for v in bijector.variables])
y = self.evaluate(bijector.forward(x))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=1,
inverse_event_ndims=1,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
# The jacobian is not yet broadcast, since it is constant.
fldj = fldj + tf.zeros(tf.shape(x)[:-1], dtype=x.dtype)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=1)
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldj(self):
softmax = tfb.SoftmaxCentered()
x = np.linspace(-15, 15, num=10).reshape(5, 2).astype(np.float64)
fldj = softmax.forward_log_det_jacobian(x, event_ndims=1)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
softmax, x, event_ndims=1)
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldjNormalCDF(self):
# b/137367959 test failure trigger case (resolved by using
# experimental_use_pfor=False as fallback instead of primary in
# bijector_test_util.get_fldj_theoretical)
bijector = tfb.TransformDiagonal(diag_bijector=tfb.NormalCDF())
x = np.zeros([0, 0])
fldj = bijector.forward_log_det_jacobian(x, event_ndims=2)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=2, inverse_event_ndims=2)
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def _checkEqualTheoreticalFldj(self, x):
"""Helper for `testJacobian`."""
event_ndims = int(self.evaluate(tf.rank(x)) - 1)
self.assertGreaterEqual(event_ndims, 1)
bijector = tfb.Cumsum(axis=-event_ndims)
fldj = bijector.forward_log_det_jacobian(
self._build_tensor(x), event_ndims=event_ndims)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=event_ndims)
fldj_, fldj_theoretical_ = self.evaluate([fldj, fldj_theoretical])
self.assertAllEqual(np.zeros_like(fldj_), fldj_)
self.assertAllClose(np.zeros_like(fldj_theoretical_), fldj_theoretical_)
def testJacobian(self):
bijector = tfb.MatrixInverseTriL()
batch_size = 5
for ndims in range(2, 5):
x_ = np.tril(
np.random.uniform(
-1., 1., size=[batch_size, ndims, ndims]).astype(np.float64))
fldj = bijector.forward_log_det_jacobian(x_, event_ndims=2)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x_, event_ndims=2,
input_to_unconstrained=tfb.Invert(tfb.FillTriangular()),
output_to_unconstrained=tfb.Invert(tfb.FillTriangular()))
fldj_, fldj_theoretical_ = self.evaluate([fldj, fldj_theoretical])
self.assertAllClose(fldj_, fldj_theoretical_)
def testTheoreticalFldj(self, data):
if not tf.executing_eagerly():
msg = ('Testing eager mode only because graph is very slow, '
'presumably due to costly graph construction.')
self.skipTest(msg)
if JAX_MODE: # TODO(b/160167257): Eliminate this workaround.
# get_fldj_theoretical uses tfp.math.batch_jacobian and assumes the
# behavior of the bijector does not vary by position. In this case, it
# can, so we must vmap the result.
batch_shape = [1]
else:
# get_fldj_theoretical test rig requires 1-d batches.
batch_shape = data.draw(tfp_hps.shapes(min_ndims=1, max_ndims=1))
hp.note('batch shape: {}'.format(batch_shape))
bijector = data.draw(rq_splines(batch_shape=batch_shape, dtype=tf.float64))
self.assertEqual(tf.float64, bijector.dtype)
bw, bh, kd = self.evaluate(
[bijector.bin_widths, bijector.bin_heights, bijector.knot_slopes])
hp.note('bw: {!r}\nbh: {!r}\nkd: {!r}'.format(bw, bh, kd))
x_shp = ((bw + bh)[..., :-1] + kd).shape[:-1]
if x_shp[-1] == 1: # Possibly broadcast the x dim.
dim = data.draw(hps.integers(min_value=1, max_value=7))
x_shp = x_shp[:-1] + (dim,)
x = np.linspace(-4.9, 4.9, np.prod(x_shp), dtype=np.float64).reshape(*x_shp)
hp.note('x: {!r}'.format(x))
y = self.evaluate(bijector.forward(x))
hp.note('x: {!r}'.format(x))
bijector_test_util.assert_bijective_and_finite(
bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=0,
inverse_event_ndims=0,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=0)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector, x, event_ndims=0)
self.assertAllClose(
self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testJacobianWithLKJSamples(self, dimension, concentration):
bijector = tfb.CorrelationCholesky()
lkj_dist = lkj.LKJ(dimension=dimension,
concentration=np.float64(concentration),
input_output_cholesky=True)
batch_size = 10
y = self.evaluate(
lkj_dist.sample([batch_size], seed=test_util.test_seed()))
x = self.evaluate(bijector.inverse(y))
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector,
x,
event_ndims=1,
inverse_event_ndims=2,
output_to_unconstrained=OutputToUnconstrained())
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
def testTheoreticalFldj(self):
bijector = tfb.CorrelationCholesky()
x = np.linspace(-50, 50, num=30).reshape(5, 6).astype(np.float64)
y = self.evaluate(bijector.forward(x))
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=1,
inverse_event_ndims=2,
rtol=1e-5)
fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
fldj_theoretical = bijector_test_util.get_fldj_theoretical(
bijector,
x,
event_ndims=1,
inverse_event_ndims=2,
output_to_unconstrained=OutputToUnconstrained())
self.assertAllClose(self.evaluate(fldj_theoretical),
self.evaluate(fldj),
atol=1e-5,
rtol=1e-5)
