Is there a stable JAX implementation for sampling from the von Mises distribution?

[benjamzc]

I'm working on a reinforcement learning project that needs to sample actions from a von Mises distribution VonMises(μ, κ) entirely in JAX. So far, I’ve tried a hand-coded accept–reject (Best & Fisher) approach, but it seems to have performance/accuracy issues (or possibly occasional infinite loops if κ is large).

I’d like a stable, well-tested sampler that stays entirely within JAX (no CPU crossing), so I can do policy gradients with log-likelihood calculations differentiable w.r.t. my network parameters. Broadly, it should work reliably for moderate κ (e.g., up to ~30).

Has anyone implemented a robust JAX version of the von Mises sampler?

Any references to existing implementations or best practices would be appreciated.

[jakevdp]

I have not seen any JAX implementation of a von Mises sampler, but perhaps you could post your approach here and we could see whether there are any obvious improvements? Are you using something like the Best & Fisher 1979 approach?