On the usage of Pystan with other sampling libraries in python

Developers
1

Hi! I was at StanCon this week and really appreciated all of the discourse and presentations. One thing that came up in a conversation with @WardBrian was the ability to benchmark other tools’ samplers with Stan-based models; for example, BlackJAX for JAX-based computation. While I think there’s a lot of meat to chew on regarding Brian’s talk about making a stan-like interface via, e.g., a lightweight interface like densityjax, I wanted to contribute my own two-cents on the problem of “we already have a well-understood stan model and want to see how X tool works on it”. The best tool I could figure out for this was a custom fork of Pystan (see here: GitHub - dannys4/pystan at change_function_interface · GitHub ), which slightly reformulates the log_prob interface a few ways.

The problem I originally ran into, and this was a bit niche so I don’t think very common, was that I needed to access both the value and gradient of log_prob for the same point, so I just made a simple wrapper for that as well. Second, a major point here is that samplers often prefer batching several evaluations simultaneously—consider, e.g., vmap transformations, and sampling from chains in parallel. An issue with using httpstan, though, is that asyncio really doesn’t like sending out thousands of separate requests!!! So I changed the interface in a way suitable for me, where you put in a numpy array and, if it’s a vector, each function does whatever it was going to do. If it’s a matrix, though, the loop over each parameter value happens inside the asyncio call instead of queueing up possibly a lot of parallel requests simultaneously.

you can see the diff here: Comparing stan-dev:main...dannys4:change_function_interface · stan-dev/pystan · GitHub . Not really much, and pretty readable, but it does slightly change the interface (and type hinting). I also created a really coarse connection example to JAX here, showing how to really poorly connect pystan to BlackJAX—I’m sure this probably exists somewhere else, but I didn’t immediately see it—and I did run into some issues that I think have to do with the HTTP requests being issued too quickly, but I’m not really sure.

Regardless, I did use my fork of pystan to pretty constructively compare between different sampling methods using posteriordb for a paper. I had a lot more luck using it with torch compared to JAX, actually, but that’s neither here nor there (torch’s approach to graph breaks is significantly more stan-friendly than JAX’s really strict rulebook).

Anyway, this may be of interest/use to people. Happy to work towards merging something in or just letting the branch sit unmaintained and stable.

2

BTW, on the use of Pystan vs. cmdstanpy, I chose Pystan to have more fine-tuned ability to access the models themselves iirc. This was like a year ago at this point, though, so please don’t interrogate me w.r.t. why choose one versus another.

3

Hi @dgsharp -

Nice meeting you! Are you aware of BridgeStan (BridgeStan – efficient, in-memory access to the methods of a Stan model — BridgeStan documentation)? It doesn’t currently have any batched signatures, but it should be more efficient and fit for this purpose than PyStan

4

I totally missed bridgestan when I was looking through these options, thanks for letting me know. I’ll definitely work it into my workflow next time I need something like this. I did find that using PyStan this way was indeed more… fragile, so this looks much nicer.

Out of curiosity, and this is mostly unrelated, does the FFI used in bridgestan somehow agree with the FFI used in JAX?

5

I don’t believe it does, since I think the Jax FFI stuff is newer than bridgestan, and we’ve also tried to avoid depending on any language-specific headers in the bindings.

So you’re stuck using the host callback mechanism as in bridgestan/python/blackjax_example.py at examples/blackjax · WardBrian/bridgestan · GitHub or GitHub - giladturok/jax-stan: Use Stan models in JAX for Bayesian inference. · GitHub, which does have some non trivial overhead. This may go down if we supported batching in the compiled code, which is a reasonable feature to add

6

Investing too much in experimental features of a google package is clearly dangerous anyway.

  1. If I’m just benchmarking methods, how much overhead do you actually expect this to be? I assume that if I needed to do performance-critical operations, it would matter, but is the overhead really that much?
  2. How much work would implementing batching be? I assume the tricky bit is now having to pass the batch size down all the way to where eigen wraps the pointer and, at that point, you’re changing a bunch of function interfaces to (at least) add a new argument for dimension, which is obviously undesirable.
7

The loop could happen anywhere once it’s in the C++ layer with roughly the same efficiency. It could also be a new signature to avoid the cascading issue you identified (though we’d probably want to just overload the behavior of the higher-level functions in e.g. Python)

The overhead isn’t killer if you’re just trying to count something like ess/grad, as I recall, but it’s been a while. @Bob_Carpenter might recall from when Gilad or Justin were using it for experiments. I believe for MCMC tasks it was fine, but for VI where they wanted huge batch evals it was painful. But that’s exactly the case we could likely improve