This is my implementation of OT-Flow using the Jax-based Equinox and Diffrax libraries. I've kept things mostly faithful to the original implementation, other than a few stylistic and package architecture changes. This is still a work in progress; questions, comments and PRs are welcome.
- Onken, D., Wu Fung, S., Li, X., & Ruthotto, L. (2021). OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport. Proceedings of the AAAI Conference on Artificial Intelligence, 35(10), 9223-9232. https://doi.org/10.1609/aaai.v35i10.17113
- Kidger, P. On Neural Differential Equations. Doctoral Dissertation. University of Oxford. 2021. URL: https://arxiv.org/abs/2202.02435. Accessed 07 March 2023.
- Kidger, P., Garcia, C., Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing Systems, 2021. URL: https://arxiv.org/abs/2111.00254. Accessed 07 March 2023.