ml-explore
[Feature] Add support mx.random.multivariate_normal()
@NNSSA and I are working on a sampling package for mlx (https://github.com/tedwards2412/samplex) and it would be extremely useful to have these two functions to do more generic sampling. The latter will involve adding more functionality to the core.linalg sub-package. Is this likely to come in a future update? Happy to help if needed!
Let's do these as two issues as diag is much easier than multivariate normal. I assume for multivariate normal you need a non-diagonal covariance?
Leaving this issue for mx.random.multivariate_normal and created #503 for diag
FYI: for multivariate normal we probably 🤔 need matrix inversion e.g. mx.linalg.inv. Which will also probably help with other things.
Great, thanks! And yes, non-diagonal covariance would be essential for this.
Cool package by the way! You should add a little quick start/usage guide (when it's ready for it).
Thanks! Will definitely do :-) Related, I think numpy uses a singular value decomposition to compute a multivariate normal.
We have a PR out for QR #310. I think SVD and Cholesky would go similarly. The main issue is there are no Metal implementations for most of Lapack so a lot of this will be CPU only until we can get some kernels implemented.
@awni We finally got round to adding the quick start you recommended on https://github.com/tedwards2412/samplex. The sampling seems to work well with mlx so far. Looking forward to working on this more in the future!
That's awesome!! Out of curiosity, could you tell me a bit more about (some) intended uses for the package? I would love to point people to it if you are ok with that and I understand a bit more in what cases you are targeting.
Overall the goal is to see if we can allow people to quickly run fairly large scale MCMC sampling locally rather than having to run on a cluster. But it's also just a research project for us to see if there are particular sampling algorithms that are substantially better when you can switch between running on CPU and GPU without any overhead; I don't think has been explored at all before.