We develop:
- GenJAX, an open-source probabilistic programming system with compilation to GPUs and TPUs.
We also publish papers on elements of our system:
- Probabilistic programming with programmable variational inference, a new design for parametric variational inference (VI) which supports programmability of VI objective functions, unbiased gradient estimators for those functions, and variational families.