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)))