def testJacobianDiagonalScaling(self, dtype):
tf_dtype = tf.as_dtype(dtype)
num_dims = 10
matrix_diagonal = np.random.uniform(size=[num_dims]).astype(dtype)
scaling_matrix = np.diag(matrix_diagonal)
one_time_scale_matrix = np.diag(np.exp(matrix_diagonal))
scale_ode_fn = lambda t, z: tf.linalg.matvec(scaling_matrix, z)
trace_augmentation_fn = tfb.ffjord.trace_jacobian_exact
bijector = tfb.FFJORD(trace_augmentation_fn=trace_augmentation_fn,
state_time_derivative_fn=scale_ode_fn,
dtype=tf_dtype)
x = np.random.uniform(size=[1, num_dims]).astype(dtype)
y = np.matmul(x, one_time_scale_matrix)
expected_forward_log_det_jacobian_value = np.log(
np.prod(np.exp(matrix_diagonal)))
expected_fldj = np.array([expected_forward_log_det_jacobian_value])
expected_ildj = np.array([-expected_forward_log_det_jacobian_value])
self.assertAllClose(y,
self.evaluate(bijector.forward(x)),
rtol=0.0,
atol=1e-3)
self.assertAllClose(x,
self.evaluate(bijector.inverse(y)),
rtol=0.0,
atol=1e-3)
self.assertAllClose(
expected_ildj,
self.evaluate(bijector.inverse_log_det_jacobian(y, event_ndims=1)),
)
self.assertAllClose(
expected_fldj,
self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)),
)
def testJacobianScaling(self, dtype):
tf_dtype = tf.as_dtype(dtype)
scaling_by_two_exp = np.log(2.0)
scale_ode_fn = lambda t, z: scaling_by_two_exp * z
trace_augmentation_fn = tfb.ffjord.trace_jacobian_exact
bijector = tfb.FFJORD(trace_augmentation_fn=trace_augmentation_fn,
state_time_derivative_fn=scale_ode_fn,
dtype=tf_dtype)
x = np.array([[0.0], [1.0]], dtype=dtype)
y = np.array([[0.0], [2.0]], dtype=dtype)
expected_forward_log_det_jacobian = np.array([np.log(2.0)] * 2,
dtype=dtype)
expected_inverse_log_det_jacobian = np.array([np.log(0.5)] * 2,
dtype=dtype)
self.assertAllClose(y,
self.evaluate(bijector.forward(x)),
rtol=0.0,
atol=1e-4)
self.assertAllClose(x,
self.evaluate(bijector.inverse(y)),
rtol=0.0,
atol=1e-4)
self.assertAllClose(
expected_inverse_log_det_jacobian,
self.evaluate(bijector.inverse_log_det_jacobian(y, event_ndims=1)),
)
self.assertAllClose(
expected_forward_log_det_jacobian,
self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)),
)
def testHutchinsonsNormalEstimator(self, dtype):
seed = 42
tf_dtype = tf.as_dtype(dtype)
num_dims = 10
np.random.seed(seed=seed)
matrix_diagonal = np.random.uniform(size=[num_dims]).astype(dtype)
scaling_matrix = np.diag(matrix_diagonal)
one_time_scale_matrix = np.diag(np.exp(matrix_diagonal))
scale_ode_fn = lambda t, z: tf.linalg.matvec(scaling_matrix, z)
def trace_augmentation_fn(ode_fn, z_shape, dtype):
return tfb.ffjord.trace_jacobian_hutchinson(ode_fn,
z_shape,
dtype,
num_samples=128,
seed=seed)
bijector = tfb.FFJORD(trace_augmentation_fn=trace_augmentation_fn,
state_time_derivative_fn=scale_ode_fn,
dtype=tf_dtype)
x = np.random.uniform(size=[1, num_dims]).astype(dtype)
y = np.matmul(x, one_time_scale_matrix)
expected_forward_log_det_jacobian_value = np.log(
np.prod(np.exp(matrix_diagonal)))
expected_fldj = np.array([expected_forward_log_det_jacobian_value])
expected_ildj = np.array([-expected_forward_log_det_jacobian_value])
self.assertAllClose(y,
self.evaluate(bijector.forward(x)),
rtol=0.0,
atol=1e-3)
self.assertAllClose(x,
self.evaluate(bijector.inverse(y)),
rtol=0.0,
atol=1e-3)
self.assertAllClose(
expected_ildj,
self.evaluate(bijector.inverse_log_det_jacobian(y, event_ndims=1)),
atol=7e-1)
self.assertAllClose(
expected_fldj,
self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)),
atol=7e-1)
def testBijectorConditionKwargs(self, dtype):
if not tf2.enabled():
self.skipTest('b/152464477')
tf_dtype = tf.as_dtype(dtype)
def conditional_ode_fn(t, z, c):
del t # unused.
return tf.ones_like(z) * c**2
trace_augmentation_fn = tfb.ffjord.trace_jacobian_exact
bijector = tfb.FFJORD(trace_augmentation_fn=trace_augmentation_fn,
state_time_derivative_fn=conditional_ode_fn,
dtype=tf_dtype)
x = tf.zeros((2, 5), dtype=tf_dtype)
y = tf.ones((2, 5), dtype=tf_dtype) * 4
c = tf.ones((2, 5), dtype=tf_dtype) * 2
expected_log_det_jacobian = np.zeros(2, dtype=dtype)
expected_dy_dc = np.ones((2, 5), dtype=dtype) * 4
def grad_fn(c):
y = bijector.forward(x, c=c)
return y
dy_dc = self.evaluate(tfp_gradient.value_and_gradient(grad_fn, c)[1])
self.assertStartsWith(bijector.name, 'ffjord')
self.assertAllClose(self.evaluate(y),
self.evaluate(bijector.forward(x, c=c)),
atol=1e-5)
self.assertAllClose(self.evaluate(x),
self.evaluate(bijector.inverse(y, c=c)),
atol=1e-5)
self.assertAllClose(
expected_log_det_jacobian,
self.evaluate(
bijector.inverse_log_det_jacobian(y, event_ndims=1, c=c)))
self.assertAllClose(
expected_log_det_jacobian,
self.evaluate(
bijector.forward_log_det_jacobian(x, event_ndims=1, c=c)))
self.assertAllClose(expected_dy_dc, dy_dc)
def testBijector(self, dtype):
tf_dtype = tf.as_dtype(dtype)
move_ode_fn = lambda t, z: tf.ones_like(z)
trace_augmentation_fn = tfb.ffjord.trace_jacobian_exact
bijector = tfb.FFJORD(trace_augmentation_fn=trace_augmentation_fn,
state_time_derivative_fn=move_ode_fn,
dtype=tf_dtype)
x = np.zeros((2, 5), dtype=dtype)
y = np.ones((2, 5), dtype=dtype)
expected_log_det_jacobian = np.zeros(2, dtype=dtype)
self.assertStartsWith(bijector.name, 'ffjord')
self.assertAllClose(y, self.evaluate(bijector.forward(x)))
self.assertAllClose(x, self.evaluate(bijector.inverse(y)))
self.assertAllClose(
expected_log_det_jacobian,
self.evaluate(bijector.inverse_log_det_jacobian(y, event_ndims=1)))
self.assertAllClose(
expected_log_det_jacobian,
self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)))