def test_poisson_likelihood(rate: float, hybridize: bool) -> None:
"""
Test to check that maximizing the likelihood recovers the parameters
"""
# generate samples
rates = mx.nd.zeros(NUM_SAMPLES) + rate
distr = Poisson(rates)
samples = distr.sample()
init_biases = [inv_softplus(rate - START_TOL_MULTIPLE * TOL * rate)]
rate_hat = maximum_likelihood_estimate_sgd(
PoissonOutput(),
samples,
init_biases=init_biases,
hybridize=hybridize,
learning_rate=PositiveFloat(0.05),
num_epochs=PositiveInt(20),
)
print("rate:", rate_hat)
assert (np.abs(rate_hat[0] - rate) < TOL *
rate), f"mu did not match: rate = {rate}, rate_hat = {rate_hat}"
alpha=0.5 * mx.nd.ones(shape=BATCH_SHAPE),
beta=0.5 * mx.nd.ones(shape=BATCH_SHAPE),
),
StudentT(
mu=mx.nd.zeros(shape=BATCH_SHAPE),
sigma=mx.nd.ones(shape=BATCH_SHAPE),
nu=mx.nd.ones(shape=BATCH_SHAPE),
),
Dirichlet(alpha=mx.nd.ones(shape=BATCH_SHAPE)),
Laplace(mu=mx.nd.zeros(shape=BATCH_SHAPE),
b=mx.nd.ones(shape=BATCH_SHAPE)),
NegativeBinomial(
mu=mx.nd.zeros(shape=BATCH_SHAPE),
alpha=mx.nd.ones(shape=BATCH_SHAPE),
),
Poisson(rate=mx.nd.ones(shape=BATCH_SHAPE)),
Uniform(
low=-mx.nd.ones(shape=BATCH_SHAPE),
high=mx.nd.ones(shape=BATCH_SHAPE),
),
PiecewiseLinear(
gamma=mx.nd.ones(shape=BATCH_SHAPE),
slopes=mx.nd.ones(shape=(3, 4, 5, 10)),
knot_spacings=mx.nd.ones(shape=(3, 4, 5, 10)) / 10,
),
MixtureDistribution(
mixture_probs=mx.nd.stack(
0.2 * mx.nd.ones(shape=BATCH_SHAPE),
0.8 * mx.nd.ones(shape=BATCH_SHAPE),
axis=-1,
),
),
(
Laplace(mu=mx.nd.zeros(shape=(3, 4, 5)),
b=mx.nd.ones(shape=(3, 4, 5))),
(3, 4, 5),
(),
),
(
NegativeBinomial(
mu=mx.nd.zeros(shape=(3, 4, 5)),
alpha=mx.nd.ones(shape=(3, 4, 5)),
),
(3, 4, 5),
(),
),
(Poisson(rate=mx.nd.ones(shape=(3, 4, 5))), (3, 4, 5), ()),
(
Uniform(
low=-mx.nd.ones(shape=(3, 4, 5)),
high=mx.nd.ones(shape=(3, 4, 5)),
),
(3, 4, 5),
(),
),
(
PiecewiseLinear(
gamma=mx.nd.ones(shape=(3, 4, 5)),
slopes=mx.nd.ones(shape=(3, 4, 5, 10)),
knot_spacings=mx.nd.ones(shape=(3, 4, 5, 10)) / 10,
),
(3, 4, 5),