SVI with Von Mises distribution

numpyro
1

Hello everybody. I am encountering a problem with the VonMises distribution, and in particular with its concentration parameter. I am trying to perform a very simple MLE of a hierarchical model.

def model(X):
    
    plate =  numpyro.plate("data", Nc)
    kappa = numpyro.param("kappa", 1., constraint = numpyro.distributions.constraints.positive)
    mu = numpyro.param("mu", 0., constraint = numpyro.distributions.constraints.interval(-np.pi, np.pi))
    
    with plate:
        phi = numpyro.sample("phi", numpyro.distributions.VonMises(mu, kappa))

    with plate:
        numpyro.sample("X", numpyro.distributions.Normal(phi, 1.), obs=X)

def guide(X):
    pass

When I run SVI, I get:

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.

Playing around with the model I noticed that if I fix the kappa (concentration parameter) to constant, and optimize only the mean of the VonMises, it works.
Also following the docs I added on top of my function:

@handlers.reparam(config={"phi": CircularReparam()})

which changes the error message to:

NotImplementedError:

I finally tried (based on this) changing my sample to:

    with plate:
        with handlers.reparam(config={'phi': CircularReparam()}):
            phi = numpyro.sample("phi", numpyro.distributions.VonMises(mu, 2.0))

Which also ends up with the NotImplementedError.

Is there a trick, or the VonMises distribution is just not well implemented yet?
Have a good day!
A

2

One option is to support forward mode differentiation in SVI (basically replace this line). Could you make an issue in github?