Implement algorithms in JAX

[x-tabdeveloping]

Rationale

Sklearn implementations of certain algorithms are not very scalable, and struggle with larger datasets and number of topics (e.g. GMM). Additionally being able to run inference on GPU and TPU would be a very nice touch.

Proposal

Implement Gaussian Mixtures (maximum likelihood - EM), NMF (coordinate descent) and FastICA in JAX, and pack it in a separate Python package that can serve as an optional dependeny.

Notes

Here are some possible implementations that could be used:

• GMM: https://gist.github.com/junpenglao/70e9bc3d60b310646bfa189394db1a29
• NMF: https://jaxopt.github.io/stable/auto_examples/constrained/nmf.html
• ICA: https://github.com/ylefay/independent_component_analysis

[x-tabdeveloping]

GMM's slowness is in-part addressed by #9.

[x-tabdeveloping]

I have also tried implementing GMM using the EM implementation I listed earlier. It did not work for high dimensionality and a lot of divergences happened. Then I tried explicitly writing out the likelihood in JAX and optimizing it with Optax and JaxOpt, both libraries resulted in divergences in under ten iterations (nans all over the place). I also tried using Bayeux with both VI, which gave some weird errors that I had nothing to do with and NUTS implementations, which were too slow to be useful and crashed my computer immediately. I unfortunately think that it is beyond the scope of the project to create and maintain a good JAX implementation of multivariate Gaussian Mixtures.

We could still try ICA and NMF, as those are probably a bit more robust to divergences, but I'm growing a bit skeptical of this idea by the day.

[x-tabdeveloping]

PyMC's VI worked, but it is incredibly slow, so we would just be back at square one basically.