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}"
(3, 4, 5), (), ), ( MixtureDistributionOutput([ MultivariateGaussianOutput(dim=5), MultivariateGaussianOutput(dim=5), ]), mx.nd.random.normal(shape=(3, 4, 10)), [None, mx.nd.ones(shape=(3, 4, 5))], [None, mx.nd.ones(shape=(3, 4, 5))], (3, 4), (5, ), ), ( PoissonOutput(), mx.nd.random.normal(shape=(3, 4, 5, 6)), [None], [None, mx.nd.ones(shape=(3, 4, 5))], (3, 4, 5), (), ), ( DeterministicOutput(42.0), mx.nd.random.normal(shape=(3, 4, 5, 6)), [None, mx.nd.ones(shape=(3, 4, 5))], [None, mx.nd.ones(shape=(3, 4, 5))], (3, 4, 5), (), ), ],