JAX-based pymdp v2 for speed improvement?

[mthiboust]

Hello,
According to this recent publication about the Cpp-AIF github project, the main painpoint of pymdp is its speed. Looking at pymdp source code, I see that it could greatly benefit from more vectorized code and hardware acceleration. With its Numpy-like API, JAX could be a great choice here. Are there any plans to rewrite this library using JAX? I'm pretty confident that it could bring several orders of magnitude faster computations while keeping the simplicity of python over C++.

[conorheins]

That is exactly our thinking as well! We have been working on the JAX backend for quite some time now actually -- see this branch : https://github.com/infer-actively/pymdp/tree/agent_jax

We haven't benchmarked it yet against that CPP paper but that would be nice to do at some point

[conorheins]

Note that this new backend allows fitting of active inference parameters to observed data (e.g. behaviour) via frameworks like numpyro, and also opens it up to hybridisation with autodifferentiable programs like deep nets

[mthiboust]

Nice, I'll have a look