mu=mx.nd.zeros(shape=(3, 4, 5)),
L=make_nd_diag(
F=mx.nd, x=mx.nd.ones(shape=(3, 4, 5)), d=5),
),
],
),
(3, 4),
(5, ),
),
(
TransformedDistribution(
StudentT(
mu=mx.nd.zeros(shape=(3, 4, 5)),
sigma=mx.nd.ones(shape=(3, 4, 5)),
nu=mx.nd.ones(shape=(3, 4, 5)),
),
[
AffineTransformation(scale=1e-1 +
mx.nd.random.uniform(shape=(3, 4, 5)))
],
),
(3, 4, 5),
(),
),
(
TransformedDistribution(
MultivariateGaussian(
mu=mx.nd.zeros(shape=(3, 4, 5)),
L=make_nd_diag(F=mx.nd, x=mx.nd.ones(shape=(3, 4, 5)),
d=5),
),
BATCH_SHAPE = (3, 4, 5)
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,)),