Exposing Stan functions in Python via PyBind

Publicity
,
1

Hi all -

I’ve recently done some experimenting in how to expose functions written in a Stan model to Python. I’ve cleaned up the results of these and put them up here: GitHub - WardBrian/pybind_expose_stan_fns: Exposing Stan functions in Python. This is modeled after the similar functionality Rok has built for cmdstanr users.

A Stan model like

functions {
  real my_log1p_exp(real x) {
    return log1p_exp(x);
  }

  real array_fun(array[] real a) {
    return sum(a);
  }

  real int_array_fun(array[] int a) {
    return sum(a);
  }

  vector my_vector_mul_by_5(vector x) {
    vector[num_elements(x)] result = x * 5.0;
    return result;
  }

  int int_only_multiplication(int a, int b) {
    return a * b;
  }

  real test_lgamma(real x) {
    return lgamma(x);
  }

  // test special functions
  void test_lp(real a) {
    a ~ normal(0, 1);
  }

  real test_rng(real a) {
    return normal_rng(a, 1);
  }

  real test_lpdf(real a, real b) {
    return normal_lpdf(a | b, 1);
  }

  void test_printing(){
    print("hi there!");
  }
}

Gets compiled to a Python module:

>>> import numpy as np
>>> import basic
>>> dir(basic)
['StanAccumulator', 'StanRNG', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__spec__', 
 'array_fun', 'int_array_fun', 'int_only_multiplication', 'my_log1p_exp', 'my_vector_mul_by_5', 'test_lgamma', 'test_lp', 'test_lpdf', 'test_printing', 'test_rng']
>>> basic.my_vector_mul_by_5(np.arange(1,5))
array([ 5., 10., 15., 20.])

This is not a fully production-ready effort, but it works for all kinds of user defined functions on Ubuntu. If someone wants to try it out on Mac I’d love to hear how successful it is or not, and I’m sure we could tweak it to work. I’m not confident this kind of thing will ever work on Windows natively, but happy to try it.

9 Likes
2

Sorry that I didn’t notice this until now!

I am playing with it on an Apple Silicon M1 Mac Studio. The first — easy to fix — problem was changing libtbb.so.2 to libtbb.dylib.

After that, I discovered that pybind11 requires -undefined dynamic_lookup on macOS, according to the docs.

But then it seems to work!

>>> import test_basic
20
[ 5. 10. 15. 20. 25. 30. 35. 40. 45.]
15.5
10.875817604889768
hi there!
Caught error: Exception: Oops! (in 'basic.stan', line 45, column 4 to column 20)
1 Like
3

Thanks for taking a look @defjaf! I just updated the repo to try to determine if you are running on MacOS and do the steps you described automatically.

4

Seems to work! I’ll try it on an intel Mac later. But this is really great – thanks.

5

Works on intel, too!.. Looking forward to using it in a real project…

1 Like
6

This turned out to be too pessimistic! I have a version of this also working on Windows, granted it requires a bit of a nonstandard setup (I’m using Microsoft’s Build Tools with the optional Clang/LLVM features)