def _testLegalInputs(self, shift=None, scale_params=None, x=None): def _powerset(x): s = list(x) return itertools.chain.from_iterable( itertools.combinations(s, r) for r in range(len(s) + 1)) for args in _powerset(scale_params.items()): with self.cached_session(): args = dict(args) scale_args = dict({"x": x}, **args) scale = self._makeScale(**scale_args) # We haven't specified enough information for the scale. if scale is None: with self.assertRaisesRegexp(ValueError, ("must be specified.")): bijector = Affine(shift=shift, **args) else: bijector = Affine(shift=shift, **args) np_x = x # For the case a vector is passed in, we need to make the shape # match the matrix for matmul to work. if x.ndim == scale.ndim - 1: np_x = np.expand_dims(x, axis=-1) forward = np.matmul(scale, np_x) + shift if x.ndim == scale.ndim - 1: forward = np.squeeze(forward, axis=-1) self.assertAllClose(forward, bijector.forward(x).eval()) backward = np.linalg.solve(scale, np_x - shift) if x.ndim == scale.ndim - 1: backward = np.squeeze(backward, axis=-1) self.assertAllClose(backward, bijector.inverse(x).eval()) scale *= np.ones(shape=x.shape[:-1], dtype=scale.dtype) ildj = -np.log(np.abs(np.linalg.det(scale))) # TODO(jvdillon): We need to make it so the scale_identity_multiplier # case does not deviate in expected shape. Fixing this will get rid of # these special cases. if (ildj.ndim > 0 and (len(scale_args) == 1 or (len(scale_args) == 2 and scale_args.get( "scale_identity_multiplier", None) is not None))): ildj = np.squeeze(ildj[0]) elif ildj.ndim < scale.ndim - 2: ildj = np.reshape(ildj, scale.shape[0:-2]) self.assertAllClose( ildj, bijector.inverse_log_det_jacobian( x, event_ndims=1).eval())
def testNoBatchMultivariateIdentity(self): with self.cached_session() as sess: placeholder = array_ops.placeholder(dtypes.float32, name="x") def static_run(fun, x, **kwargs): return fun(x, **kwargs).eval() def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) return sess.run(fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = [1., -1] # Multivariate # Corresponds to scale = [[1., 0], [0, 1.]] bijector = Affine(shift=mu) x = [1., 1] # matmul(sigma, x) + shift # = [-1, -1] + [1, -1] self.assertAllClose([2., 0], run(bijector.forward, x)) self.assertAllClose([0., 2], run(bijector.inverse, x)) # x is a 2-batch of 2-vectors. # The first vector is [1, 1], the second is [-1, -1]. # Each undergoes matmul(sigma, x) + shift. x = [[1., 1], [-1., -1]] self.assertAllClose([[2., 0], [0., -2]], run(bijector.forward, x)) self.assertAllClose([[0., 2], [-2., 0]], run(bijector.inverse, x)) self.assertAllClose( 0., run(bijector.inverse_log_det_jacobian, x, event_ndims=1))
def testCompareToBijector(self): """Demonstrates equivalence between TD, Bijector approach and AR dist.""" sample_shape = np.int32([4, 5]) batch_shape = np.int32([]) event_size = np.int32(2) with self.cached_session() as sess: batch_event_shape = np.concatenate([batch_shape, [event_size]], axis=0) sample0 = array_ops.zeros(batch_event_shape) affine = Affine(scale_tril=self._random_scale_tril(event_size)) ar = autoregressive_lib.Autoregressive(self._normal_fn(affine), sample0, validate_args=True) ar_flow = MaskedAutoregressiveFlow(is_constant_jacobian=True, shift_and_log_scale_fn=lambda x: [None, affine.forward(x)], validate_args=True) td = transformed_distribution_lib.TransformedDistribution( distribution=normal_lib.Normal(loc=0., scale=1.), bijector=ar_flow, event_shape=[event_size], batch_shape=batch_shape, validate_args=True) x_shape = np.concatenate([sample_shape, batch_shape, [event_size]], axis=0) x = 2. * self._rng.random_sample(x_shape).astype(np.float32) - 1. td_log_prob_, ar_log_prob_ = sess.run( [td.log_prob(x), ar.log_prob(x)]) self.assertAllClose(td_log_prob_, ar_log_prob_, atol=0., rtol=1e-6)
def testIdentityAndDiagWithTriL(self): with self.cached_session() as sess: placeholder = array_ops.placeholder(dtypes.float32, name="x") def static_run(fun, x, **kwargs): return fun(x, **kwargs).eval() def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) return sess.run(fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. # scale = [[3., 0], [2, 4]] bijector = Affine(shift=mu, scale_identity_multiplier=1.0, scale_diag=[1., 2.], scale_tril=[[1., 0], [2., 1]]) x = [[1., 2]] # One multivariate sample. self.assertAllClose([[2., 9]], run(bijector.forward, x)) self.assertAllClose([[2 / 3., 5 / 12.]], run(bijector.inverse, x)) self.assertAllClose( -np.log(12.), run(bijector.inverse_log_det_jacobian, x, event_ndims=1))
def testBatchMultivariateFullDynamic(self): with self.cached_session() as sess: x = array_ops.placeholder(dtypes.float32, name="x") mu = array_ops.placeholder(dtypes.float32, name="mu") scale_diag = array_ops.placeholder(dtypes.float32, name="scale_diag") x_value = np.array([[[1., 1]]], dtype=np.float32) mu_value = np.array([[1., -1]], dtype=np.float32) scale_diag_value = np.array([[2., 2]], dtype=np.float32) feed_dict = { x: x_value, mu: mu_value, scale_diag: scale_diag_value, } bijector = Affine(shift=mu, scale_diag=scale_diag) self.assertAllClose([[[3., 1]]], sess.run(bijector.forward(x), feed_dict)) self.assertAllClose([[[0., 1]]], sess.run(bijector.inverse(x), feed_dict)) self.assertAllClose( [-np.log(4)], sess.run(bijector.inverse_log_det_jacobian(x, event_ndims=1), feed_dict))
def testMinEventNdimsShapeChangingRemoveDims(self): chain = Chain([ShapeChanging(3, 0)]) self.assertEqual(3, chain.forward_min_event_ndims) self.assertEqual(0, chain.inverse_min_event_ndims) chain = Chain([ShapeChanging(3, 0), Affine()]) self.assertEqual(3, chain.forward_min_event_ndims) self.assertEqual(0, chain.inverse_min_event_ndims) chain = Chain([Affine(), ShapeChanging(3, 0)]) self.assertEqual(4, chain.forward_min_event_ndims) self.assertEqual(1, chain.inverse_min_event_ndims) chain = Chain([ShapeChanging(3, 0), ShapeChanging(3, 0)]) self.assertEqual(6, chain.forward_min_event_ndims) self.assertEqual(0, chain.inverse_min_event_ndims)
def testNoBatchMultivariateRaisesWhenSingular(self): with self.cached_session(): mu = [1., -1] bijector = Affine( shift=mu, # Has zero on the diagonal. scale_diag=[0., 1], validate_args=True) with self.assertRaisesOpError("diagonal part must be non-zero"): bijector.forward([1., 1.]).eval()
def testTriLWithVDVTUpdateNoDiagonal(self): with self.cached_session() as sess: placeholder = array_ops.placeholder(dtypes.float32, name="x") def static_run(fun, x, **kwargs): return fun(x, **kwargs).eval() def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) return sess.run(fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. # Corresponds to scale = [[6, 0, 0], [1, 3, 0], [2, 3, 5]] bijector = Affine(shift=mu, scale_tril=[[2., 0, 0], [1, 3, 0], [2, 3, 4]], scale_perturb_diag=None, scale_perturb_factor=[[2., 0], [0., 0], [0, 1]]) bijector_ref = Affine(shift=mu, scale_tril=[[6., 0, 0], [1, 3, 0], [2, 3, 5]]) x = [1., 2, 3] # Vector. self.assertAllClose([5., 6, 22], run(bijector.forward, x)) self.assertAllClose(run(bijector_ref.forward, x), run(bijector.forward, x)) self.assertAllClose([1 / 3., 8 / 9., 4 / 30.], run(bijector.inverse, x)) self.assertAllClose(run(bijector_ref.inverse, x), run(bijector.inverse, x)) self.assertAllClose( -np.log(90.), run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) self.assertAllClose( run(bijector.inverse_log_det_jacobian, x, event_ndims=1), run(bijector_ref.inverse_log_det_jacobian, x, event_ndims=1))
def testIdentityWithVDVTUpdate(self): with self.cached_session() as sess: placeholder = array_ops.placeholder(dtypes.float32, name="x") def static_run(fun, x, **kwargs): return fun(x, **kwargs).eval() def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) return sess.run(fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. # Corresponds to scale = [[10, 0, 0], [0, 2, 0], [0, 0, 3]] bijector = Affine(shift=mu, scale_identity_multiplier=2., scale_perturb_diag=[2., 1], scale_perturb_factor=[[2., 0], [0., 0], [0, 1]]) bijector_ref = Affine(shift=mu, scale_diag=[10., 2, 3]) x = [1., 2, 3] # Vector. self.assertAllClose([9., 3, 8], run(bijector.forward, x)) self.assertAllClose(run(bijector_ref.forward, x), run(bijector.forward, x)) self.assertAllClose([0.2, 1.5, 4 / 3.], run(bijector.inverse, x)) self.assertAllClose(run(bijector_ref.inverse, x), run(bijector.inverse, x)) self.assertAllClose( -np.log(60.), run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) self.assertAllClose( run(bijector.inverse_log_det_jacobian, x, event_ndims=1), run(bijector_ref.inverse_log_det_jacobian, x, event_ndims=1))
def testChainAffineExp(self): scale_diag = np.array([1., 2., 3.], dtype=np.float32) chain = Chain([Affine(scale_diag=scale_diag), Exp()]) x = [0., np.log(2., dtype=np.float32), np.log(3., dtype=np.float32)] y = [1., 4., 9.] self.assertAllClose(y, self.evaluate(chain.forward(x))) self.assertAllClose(x, self.evaluate(chain.inverse(y))) self.assertAllClose( np.log(6, dtype=np.float32) + np.sum(x), self.evaluate(chain.forward_log_det_jacobian(x, event_ndims=1))) self.assertAllClose( -np.log(6, dtype=np.float32) - np.sum(x), self.evaluate(chain.inverse_log_det_jacobian(y, event_ndims=1)))
def testSampleAndLogProbConsistency(self): batch_shape = [] event_size = 2 with self.cached_session() as sess: batch_event_shape = np.concatenate([batch_shape, [event_size]], axis=0) sample0 = array_ops.zeros(batch_event_shape) affine = Affine(scale_tril=self._random_scale_tril(event_size)) ar = autoregressive_lib.Autoregressive(self._normal_fn(affine), sample0, validate_args=True) self.run_test_sample_consistent_log_prob(sess.run, ar, radius=1., center=0., rtol=0.01)
def testMinEventNdimsChain(self): chain = Chain([Exp(), Exp(), Exp()]) self.assertEqual(0, chain.forward_min_event_ndims) self.assertEqual(0, chain.inverse_min_event_ndims) chain = Chain([Affine(), Affine(), Affine()]) self.assertEqual(1, chain.forward_min_event_ndims) self.assertEqual(1, chain.inverse_min_event_ndims) chain = Chain([Exp(), Affine()]) self.assertEqual(1, chain.forward_min_event_ndims) self.assertEqual(1, chain.inverse_min_event_ndims) chain = Chain([Affine(), Exp()]) self.assertEqual(1, chain.forward_min_event_ndims) self.assertEqual(1, chain.inverse_min_event_ndims) chain = Chain([Affine(), Exp(), Softplus(), Affine()]) self.assertEqual(1, chain.forward_min_event_ndims) self.assertEqual(1, chain.inverse_min_event_ndims)
def testBatchMultivariateIdentity(self): with self.cached_session() as sess: placeholder = array_ops.placeholder(dtypes.float32, name="x") def static_run(fun, x, **kwargs): return fun(x, **kwargs).eval() def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) return sess.run(fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = [[1., -1]] # Corresponds to 1 2x2 matrix, with twos on the diagonal. scale = 2. bijector = Affine(shift=mu, scale_identity_multiplier=scale) x = [[[1., 1]]] self.assertAllClose([[[3., 1]]], run(bijector.forward, x)) self.assertAllClose([[[0., 1]]], run(bijector.inverse, x)) self.assertAllClose( -np.log(4), run(bijector.inverse_log_det_jacobian, x, event_ndims=1))
def testIdentityWithDiagUpdate(self): with self.cached_session() as sess: placeholder = array_ops.placeholder(dtypes.float32, name="x") def static_run(fun, x, **kwargs): return fun(x, **kwargs).eval() def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) return sess.run(fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. # Corresponds to scale = 2 bijector = Affine(shift=mu, scale_identity_multiplier=1., scale_diag=[1., 1., 1.]) x = [1., 2, 3] # Three scalar samples (no batches). self.assertAllClose([1., 3, 5], run(bijector.forward, x)) self.assertAllClose([1., 1.5, 2.], run(bijector.inverse, x)) self.assertAllClose( -np.log(2.**3), run(bijector.inverse_log_det_jacobian, x, event_ndims=1))
def testProperties(self): with self.cached_session(): mu = -1. # scale corresponds to 1. bijector = Affine(shift=mu) self.assertEqual("affine", bijector.name)