def test_transformed_distribution() -> None:
zero = nd.zeros(1)
one = nd.ones(1)
# If Y = -log(U) with U ~ Uniform(0, 1), then Y ~ Exponential(1)
exponential = TransformedDistribution(
Uniform(zero, one),
bijection.log,
bijection.AffineTransformation(scale=-1 * one),
)
# For Y ~ Exponential(1), P(Y) = e^{-x) ==> log P(Y) = -x
assert exponential.log_prob(1 * one).asscalar() == -1.0
assert exponential.log_prob(2 * one).asscalar() == -2.0
# If Y ~ Exponential(1), then U = 1 - e^{-Y} has Uniform(0, 1) distribution
uniform = TransformedDistribution(
exponential,
bijection.AffineTransformation(scale=-1 * one),
bijection.log.inverse_bijection(), # == bijection.exp
bijection.AffineTransformation(loc=one, scale=-1 * one),
)
# For U ~ Uniform(0, 1), log P(U) = 0
assert uniform.log_prob(0.5 * one).asscalar() == 0
assert uniform.log_prob(0.2 * one).asscalar() == 0
def distribution(self, distr_args, loc=None, scale=None, dim=None):
dist = LowrankMultivariateGaussian(dim, self.rank, *distr_args)
if loc is None and scale is None:
return dist
else:
return TransformedDistribution(
dist, [bijection.AffineTransformation(loc=loc, scale=scale)])
def distribution(self, distr_args, scale=None, **kwargs) -> Distribution:
# todo dirty way of calling for now, this can be cleaned
distr = LowrankMultivariateGaussian(self.dim, self.rank, *distr_args)
if scale is None:
return distr
else:
return TransformedDistribution(
distr, bijection.AffineTransformation(scale=scale))
def distribution(self, distr_args, scale=None, **kwargs) -> Distribution:
distr = LowrankMultivariateGaussian(self.dim, self.rank, *distr_args)
if scale is None:
return distr
else:
return TransformedDistribution(
distr, [bijection.AffineTransformation(scale=scale)]
)
def test_transformed_distribution(serialize_fn) -> None:
zero = nd.zeros(1)
one = nd.ones(1)
# If Y = -log(U) with U ~ Uniform(0, 1), then Y ~ Exponential(1)
exponential = TransformedDistribution(
Uniform(zero, one),
[bijection.log,
bijection.AffineTransformation(scale=-1 * one)],
)
exponential = serialize_fn(exponential)
# For Y ~ Exponential(1), P(Y) = e^{-x) ==> log P(Y) = -x
assert exponential.log_prob(1 * one).asscalar() == -1.0
assert exponential.log_prob(2 * one).asscalar() == -2.0
v = np.linspace(0, 5, 101)
assert np.allclose(exponential.cdf(nd.array(v)).asnumpy(), exp_cdf(v))
level = np.linspace(1.0e-5, 1.0 - 1.0e-5, 101)
qs_calc = exponential.quantile(nd.array(level)).asnumpy()[:, 0]
qs_theo = exp_quantile(level)
assert np.allclose(qs_calc, qs_theo, atol=1.0e-2)
# If Y ~ Exponential(1), then U = 1 - e^{-Y} has Uniform(0, 1) distribution
uniform = TransformedDistribution(
exponential,
[
bijection.AffineTransformation(scale=-1 * one),
bijection.log.inverse_bijection(), # == bijection.exp
bijection.AffineTransformation(loc=one, scale=-1 * one),
],
)
uniform = serialize_fn(uniform)
# For U ~ Uniform(0, 1), log P(U) = 0
assert uniform.log_prob(0.5 * one).asscalar() == 0
assert uniform.log_prob(0.2 * one).asscalar() == 0
v = np.linspace(0, 1, 101)
assert np.allclose(uniform.cdf(nd.array(v)).asnumpy(), v)
qs_calc = uniform.quantile(nd.array(level)).asnumpy()[:, 0]
assert np.allclose(qs_calc, level, atol=1.0e-2)
DISTRIBUTIONS_WITH_QUANTILE_FUNCTION = (Gaussian, Uniform, Laplace, Binned)
@pytest.mark.parametrize(
"distr",
[
TransformedDistribution(
Gaussian(
mu=mx.nd.random.uniform(shape=BATCH_SHAPE),
sigma=mx.nd.ones(shape=BATCH_SHAPE),
),
[
bij.AffineTransformation(
scale=1e-1 + mx.nd.random.uniform(shape=BATCH_SHAPE)
),
bij.softrelu,
],
),
Binned(
bin_log_probs=mx.nd.uniform(shape=BATCH_SHAPE + (23,)),
bin_centers=mx.nd.array(np.logspace(-1, 1, 23))
+ mx.nd.zeros(BATCH_SHAPE + (23,)),
),
TransformedDistribution(
Binned(
bin_log_probs=mx.nd.uniform(shape=BATCH_SHAPE + (23,)),
bin_centers=mx.nd.array(np.logspace(-1, 1, 23))
+ mx.nd.zeros(BATCH_SHAPE + (23,)),
),
