ONNX support for Stan models

[bparbhu]

Hey all,

I've been thinking more about how to improve interoperability for Stan and how Stan users can take advantage of additional performance gains by possibly switching over stan models to be saved in the onnx format. I'm interested in helping with this and would want to come up with a design doc around this. Any feedback on this is much appreciated and always happy to help.

Thanks again,

-Brian

[Craigacp]

I've thought about this in the context of other probabilistic programming systems, and I think ONNX doesn't necessarily have sufficient operators to implement MCMC, and it might not have enough for the probability distributions in general. Are you thinking about running a chain in the ONNX model, or about using a MAP estimate from a completed Bayesian inference step to then perform inference on new points. The latter seems easier, but I'm not sure if the full space of Stan models could be mapped down. For example I've previously done work on inference in topic models (which I know is not Stan's strength) and ONNX doesn't support things like the gamma function that you need to work with Dirichlet and Beta distributions.

[bparbhu]

So I have heard something similar from someone on the Stan development team as well. I think we would ideally want to have the ability to do both things with a Stan model but if you have limited functionality with the saved Stan model, then I think this wouldn't be a worthwhile effort. I always wondered why Onnx never supported LDA models as well, so thanks for letting me know about that. I actually worked on a project where we were trying to get that functionality for Onnx models as part of a chatbot.

[Craigacp]

I think this might be an interesting topic to bring up at the ONNX community meeting (https://community.linuxfoundation.org/events/details/lfhq-lf-ai-data-foundation-presents-lf-ai-data-day-onnx-community-meetup-silicon-valley/) which happens in June to see if there are other people interested. Maybe suggest probabilistic model support as a potential round table topic? Adding proper support would require new ONNX operators which is possible, but I think requires a number of people to want them.

[bparbhu]

Will do @Craigacp

[bparbhu]

I started working on something for this here https://github.com/bparbhu/onnx-stan . I think what will help us out is to leverage the work done here as well https://github.com/deepppl/stan-num-pyro . We could probably leverage this work into supporting other frameworks like pyro, numpyro, and tensorflow probability.

[bparbhu]

We'll start with stan and then we'll look at the other frameworks to see how much of the work applies. I think it will.