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