Some chat: Pyro/tf.prob are all more of an operations/stat research oriented tools, I gather but if there are some MC abstractions it might be good choice.

Yes, tested and used Jax for some prototyping but it felt like still has "two language problem". But Taichi looks quite interesting, will check out.

Of course, forgot to mention, it has to be lightweight, torch and tf are now huge platforms. Not sure, probably Julia-lang is much small.