Using non-JAX-traceable logpdfs with BlackJAX

[sethaxen]

Is it possible to use functions that are not traceable by JAX with BlackJAX? Some use cases:

• A user has some object written in another language but wrapped by Python that computes a logpdf and its gradient and wants to sample using BlackJAX.
• Similarly, a user has a logpdf function written in some other framework like PyTorch that can also compute gradients and wants to sample using BlackJAX

In the first case, the object itself might live outside of the logpdf function but be updated internally by the logpdf function using the sampled parameters in order to compute the logpdf and its gradient. This at least puts constrains on parallelism.

Is this currently possible using BlackJAX? This might be more of a JAX question than a BlackJAX one.

cc @mattjj

[rlouf]

Thanks for using blackjax! I'll let @mattjj answer the more general JAX question (you're actually not the first to ask).

On the blackjax side, there is currently no way for users to provide their own gradient because of the way the symplectic integrators are implemented. However we are currently integrating the stochastic gradient methods from SGMCMCJAX; those require the user to pass a gradient estimator and we'll harmonize the APIs so users can also provide gradients for HMC/NUTS.

[IvanYashchuk]

Yes, it seems like it's not difficult now to add support of arbitrary Python functions to work with jax.jit using jax.experimental.host_callback.
Previously, to make it work XLA CustomCall needed to be used. There's an example for using NumPy functions in the link. Once this "host" function is registered it's possible to define custom derivative rules for it following this tutorial.
The post by @dfm might be useful as well: https://dfm.io/posts/extending-jax/.

[rlouf]

@sethaxen did you find a way in the end?

[sethaxen]

No, sorry to drop the ball on this, I got overwhelmed with other things. I'll revisit this in the new year.

[rlouf]

I'm positive this is possible following the links @IvanYashchuk provided and I am thinking about writing a "How to?" tutorial. Maybe using PyTorch?

[rlouf]

It is possible, we just pushed an example in the documentation: https://blackjax-devs.github.io/blackjax/examples/howto_other_frameworks.html