  We present Rieoptax, an open source Python library for Riemannian
optimization in JAX. We show that many differential geometric primitives, such
as Riemannian exponential and logarithm maps, are usually faster in Rieoptax
than existing frameworks in Python, both on CPU and GPU. We support various
range of basic and advanced stochastic optimization solvers like Riemannian
stochastic gradient, stochastic variance reduction, and adaptive gradient
methods. A distinguishing feature of the proposed toolbox is that we also
support differentially private optimization on Riemannian manifolds.
