def testBijectiveAndFinite(self):
bijector = tfb.NormalCDF(validate_args=True)
x = np.linspace(-10., 10., num=10).astype(np.float32)
y = np.linspace(0.1, 0.9, num=10).astype(np.float32)
bijector_test_util.assert_bijective_and_finite(bijector,
x,
y,
eval_func=self.evaluate,
event_ndims=0,
rtol=1e-4)
def disabled_testFailureCase(self):
# TODO(b/140229057): This test should pass.
dist = tfd.Chi(df=np.float32(27.744131))
dist = tfd.TransformedDistribution(bijector=tfb.NormalCDF(),
distribution=dist,
batch_shape=[4])
dist = tfb.Expm1()(dist)
samps = 1.7182817 + tf.zeros_like(dist.sample())
self.assertAllClose(
dist.log_prob(samps)[0], dist[0].log_prob(samps[0]))
def disabled_testFailureCase(self): # pylint: disable=invalid-name
# TODO(b/140229057): This test should pass.
dist = tfd.Chi(df=np.float32(27.744131) *
np.ones((4, )).astype(np.float32))
dist = tfd.TransformedDistribution(bijector=tfb.NormalCDF(),
distribution=dist)
dist = tfb.Expm1()(dist)
samps = 1.7182817 + tf.zeros_like(
dist.sample(seed=test_util.test_seed()))
self.assertAllClose(
dist.log_prob(samps)[0], dist[0].log_prob(samps[0]))
def testValidateArgs(self):
bijector = tfb.NormalCDF(validate_args=True)
with self.assertRaisesOpError("must be greater than 0"):
self.evaluate(bijector.inverse(-1.))
with self.assertRaisesOpError("must be greater than 0"):
self.evaluate(bijector.inverse_log_det_jacobian(-1., event_ndims=0))
with self.assertRaisesOpError("must be less than or equal to 1"):
self.evaluate(bijector.inverse(2.))
with self.assertRaisesOpError("must be less than or equal to 1"):
self.evaluate(bijector.inverse_log_det_jacobian(2., event_ndims=0))
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 testBijector(self):
bijector = tfb.NormalCDF(validate_args=True)
self.assertEqual("normal", bijector.name)
x = np.array([[[-3.], [0.], [0.5], [4.2], [5.]]], dtype=np.float64)
# Normal distribution
normal_dist = stats.norm(loc=0., scale=1.)
y = normal_dist.cdf(x).astype(np.float64)
self.assertAllClose(y, self.evaluate(bijector.forward(x)))
self.assertAllClose(x, self.evaluate(bijector.inverse(y)), rtol=1e-4)
self.assertAllClose(
np.squeeze(normal_dist.logpdf(x), axis=-1),
self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)))
self.assertAllClose(
self.evaluate(-bijector.inverse_log_det_jacobian(y, event_ndims=1)),
self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)),
rtol=1e-4)
def testScalarCongruency(self):
bijector_test_util.assert_scalar_congruency(tfb.NormalCDF(),
lower_x=0.,
upper_x=1.,
eval_func=self.evaluate,
rtol=0.02)
def make_distribution_bijector(distribution,
name='make_distribution_bijector'):
"""Builds a bijector to approximately transform `N(0, 1)` into `distribution`.
This represents a distribution as a bijector that
transforms a (multivariate) standard normal distribution into the distribution
of interest.
Args:
distribution: A `tfd.Distribution` instance; this may be a joint
distribution.
name: Python `str` name for ops created by this method.
Returns:
distribution_bijector: a `tfb.Bijector` instance such that
`distribution_bijector(tfd.Normal(0., 1.))` is approximately equivalent
to `distribution`.
#### Examples
This method may be used to convert structured variational distributions into
MCMC preconditioners. Consider a model containing
[funnel geometry](https://crackedbassoon.com/writing/funneling), which may
be difficult for an MCMC algorithm to sample directly.
```python
model_with_funnel = tfd.JointDistributionSequentialAutoBatched([
tfd.Normal(loc=-1., scale=2., name='z'),
lambda z: tfd.Normal(loc=[0., 0., 0.], scale=tf.exp(z), name='x'),
lambda x: tfd.Poisson(log_rate=x, name='y')])
pinned_model = tfp.experimental.distributions.JointDistributionPinned(
model_with_funnel, y=[1, 3, 0])
```
We can approximate the posterior in this model using a structured variational
surrogate distribution, which will capture the funnel geometry, but cannot
exactly represent the (non-Gaussian) posterior.
```python
# Build and fit a structured surrogate posterior distribution.
surrogate_posterior = tfp.experimental.vi.build_asvi_surrogate_posterior(
pinned_model)
_ = tfp.vi.fit_surrogate_posterior(pinned_model.unnormalized_log_prob,
surrogate_posterior=surrogate_posterior,
optimizer=tf.optimizers.Adam(0.01),
num_steps=200)
```
Creating a preconditioning bijector allows us to obtain higher-quality
posterior samples, without any Gaussianity assumption, by using the surrogate
to guide an MCMC sampler.
```python
surrogate_posterior_bijector = (
tfp.experimental.bijectors.make_distribution_bijector(surrogate_posterior))
samples, _ = tfp.mcmc.sample_chain(
kernel=tfp.mcmc.DualAveragingStepSizeAdaptation(
tfp.mcmc.TransformedTransitionKernel(
tfp.mcmc.NoUTurnSampler(pinned_model.unnormalized_log_prob,
step_size=0.1),
bijector=surrogate_posterior_bijector),
num_adaptation_steps=80),
current_state=surrogate_posterior.sample(),
num_burnin_steps=100,
trace_fn=lambda _0, _1: [],
num_results=500)
```
#### Mathematical details
The bijectors returned by this method generally follow the following
principles, although the specific bijectors returned may vary without notice.
Normal distributions are reparameterized by a location-scale transform.
```python
b = tfp.experimental.bijectors.make_distribution_bijector(
tfd.Normal(loc=10., scale=5.))
# ==> tfb.Shift(10.)(tfb.Scale(5.)))
b = tfp.experimental.bijectors.make_distribution_bijector(
tfd.MultivariateNormalTriL(loc=loc, scale_tril=scale_tril))
# ==> tfb.Shift(loc)(tfb.ScaleMatvecTriL(scale_tril))
```
The distribution's `quantile` function is used, when available:
```python
d = tfd.Cauchy(loc=loc, scale=scale)
b = tfp.experimental.bijectors.make_distribution_bijector(d)
# ==> tfb.Inline(forward_fn=d.quantile, inverse_fn=d.cdf)(tfb.NormalCDF())
```
Otherwise, a quantile function is derived by inverting the CDF:
```python
d = tfd.Gamma(concentration=alpha, rate=beta)
b = tfp.experimental.bijectors.make_distribution_bijector(d)
# ==> tfb.Invert(
# tfp.experimental.bijectors.ScalarFunctionWithInferredInverse(fn=d.cdf))(
# tfb.NormalCDF())
```
Transformed distributions are represented by chaining the
transforming bijector with a preconditioning bijector for the base
distribution:
```python
b = tfp.experimental.bijectors.make_distribution_bijector(
tfb.Exp(tfd.Normal(loc=10., scale=5.)))
# ==> tfb.Exp(tfb.Shift(10.)(tfb.Scale(5.)))
```
Joint distributions are represented by a joint bijector, which converts each
component distribution to a bijector with parameters conditioned
on the previous variables in the model. The joint bijector's inputs and
outputs follow the structure of the joint distribution.
```python
jd = tfd.JointDistributionNamed(
{'a': tfd.InverseGamma(concentration=2., scale=1.),
'b': lambda a: tfd.Normal(loc=3., scale=tf.sqrt(a))})
b = tfp.experimental.bijectors.make_distribution_bijector(jd)
whitened_jd = tfb.Invert(b)(jd)
x = whitened_jd.sample()
# x <=> {'a': tfd.Normal(0., 1.).sample(), 'b': tfd.Normal(0., 1.).sample()}
```
"""
with tf.name_scope(name):
event_space_bijector = (
distribution.experimental_default_event_space_bijector())
if event_space_bijector is None: # Fail if the distribution is discrete.
raise NotImplementedError(
'Cannot transform distribution {} to a standard normal '
'distribution.'.format(distribution))
# Recurse over joint distributions.
if isinstance(distribution, joint_distribution.JointDistribution):
return joint_distribution._DefaultJointBijector( # pylint: disable=protected-access
distribution,
bijector_fn=make_distribution_bijector)
# Recurse through transformed distributions.
if isinstance(distribution,
transformed_distribution.TransformedDistribution):
return distribution.bijector(
make_distribution_bijector(distribution.distribution))
# If we've annotated a specific bijector for this distribution, use that.
if isinstance(distribution, tuple(preconditioning_bijector_fns)):
return preconditioning_bijector_fns[type(distribution)](
distribution)
# Otherwise, if this distribution implements a CDF and inverse CDF, build
# a bijector from those.
implements_cdf = False
implements_quantile = False
input_spec = tf.zeros(shape=distribution.event_shape,
dtype=distribution.dtype)
try:
callable_util.get_output_spec(distribution.cdf, input_spec)
implements_cdf = True
except NotImplementedError:
pass
try:
callable_util.get_output_spec(distribution.quantile, input_spec)
implements_quantile = True
except NotImplementedError:
pass
if implements_cdf and implements_quantile:
# This path will only trigger for scalar distributions, since multivariate
# distributions have non-invertible CDF and so cannot define a `quantile`.
return tfb.Inline(forward_fn=distribution.quantile,
inverse_fn=distribution.cdf,
forward_min_event_ndims=ps.rank_from_shape(
distribution.event_shape_tensor,
distribution.event_shape))(tfb.NormalCDF())
# If the events are scalar, try to invert the CDF numerically.
if implements_cdf and tf.get_static_value(
distribution.is_scalar_event()):
return tfb.Invert(
scalar_function_with_inferred_inverse.
ScalarFunctionWithInferredInverse(
distribution.cdf,
domain_constraint_fn=(event_space_bijector)))(
tfb.NormalCDF())
raise NotImplementedError('Could not automatically construct a '
'bijector for distribution type '
'{}; it does not implement an invertible '
'CDF.'.format(distribution))
lambda d: markov_chain._MarkovChainBijector(
chain=d,
transition_bijector=make_distribution_bijector(
d.transition_fn(
0, d.initial_state_prior.sample(seed=samplers.zeros_seed()))),
bijector_fn=make_distribution_bijector),
normal.Normal:
lambda d: tfb.Shift(d.loc)(tfb.Scale(d.scale)),
sample.Sample:
lambda d: sample._DefaultSampleBijector(
distribution=d.distribution,
sample_shape=d.sample_shape,
sum_fn=d._sum_fn(),
bijector=make_distribution_bijector(d.distribution)),
uniform.Uniform:
lambda d: (tfb.Shift(d.low)(tfb.Scale(d.high - d.low)(tfb.NormalCDF())))
}
# pylint: enable=g-long-lambda,protected-access
def make_distribution_bijector(distribution,
name='make_distribution_bijector'):
"""Builds a bijector to approximately transform `N(0, 1)` into `distribution`.
This represents a distribution as a bijector that
transforms a (multivariate) standard normal distribution into the distribution
of interest.
Args:
distribution: A `tfd.Distribution` instance; this may be a joint
distribution.
