def test_logit_normal_likelihood(mu: float, sigma: float, hybridize: bool):
"""
Test to check that maximizing the likelihood recovers the parameters
"""
# generate samples
mus = mx.nd.zeros((NUM_SAMPLES, )) + mu
sigmas = mx.nd.zeros((NUM_SAMPLES, )) + sigma
distr = LogitNormal(mus, sigmas)
samples = distr.sample()
init_biases = [
mu - START_TOL_MULTIPLE * TOL * mu,
inv_softplus(sigma - START_TOL_MULTIPLE * TOL * sigma),
]
mu_hat, sigma_hat = maximum_likelihood_estimate_sgd(
LogitNormalOutput(),
samples,
init_biases=init_biases,
hybridize=hybridize,
learning_rate=PositiveFloat(0.001),
num_epochs=PositiveInt(5),
)
assert (np.abs(mu_hat - mu) <
TOL * mu), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
assert (np.abs(sigma_hat - sigma) < TOL * sigma
), f"sigma did not match: sigma = {sigma}, sigma_hat = {sigma_hat}"
@pytest.mark.parametrize(
"distr_output",
[
BetaOutput(),
CategoricalOutput(num_cats=3),
DeterministicOutput(value=42.0),
DirichletMultinomialOutput(dim=3, n_trials=5),
DirichletOutput(dim=4),
EmpiricalDistributionOutput(num_samples=10,
distr_output=GaussianOutput()),
GammaOutput(),
GaussianOutput(),
GenParetoOutput(),
LaplaceOutput(),
LogitNormalOutput(),
LoglogisticOutput(),
LowrankMultivariateGaussianOutput(dim=5, rank=2),
MultivariateGaussianOutput(dim=4),
NegativeBinomialOutput(),
OneInflatedBetaOutput(),
PiecewiseLinearOutput(num_pieces=10),
PoissonOutput(),
StudentTOutput(),
UniformOutput(),
WeibullOutput(),
ZeroAndOneInflatedBetaOutput(),
ZeroInflatedBetaOutput(),
ZeroInflatedNegativeBinomialOutput(),
ZeroInflatedPoissonOutput(),
],