# Rieoptax: Riemannian Optimization in JAX

arXiv: https://arxiv.org/abs/2210.04840 · HTML: https://mlsanity.com/paper/2210.04840/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 116 references that are in the map.
- First posted on arXiv 10 Oct 2022.

## Builds on (most cited first, 30 of 56)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Deep Learning with Differential Privacy](https://mlsanity.com/paper/1607.00133/) (2016)
- [Membership Inference Attacks against Machine Learning Models](https://mlsanity.com/paper/1610.05820/) (2016)
- [Deep Leakage from Gradients](https://mlsanity.com/paper/1906.08935/) (2019)
- [The Secret Sharer: Evaluating and Testing Unintended Memorization in Neural Networks](https://mlsanity.com/paper/1802.08232/) (2018)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Poincar'e Embeddings for Learning Hierarchical Representations](https://mlsanity.com/paper/1705.08039/) (2017)
- [TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems](https://mlsanity.com/paper/1603.04467/) (2016)
- [Renyi Differential Privacy](https://mlsanity.com/paper/1702.07476/) (2017)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Hyperbolic Neural Networks](https://mlsanity.com/paper/1805.09112/) (2018)
- [RAPPOR: Randomized Aggregatable Privacy-Preserving Ordinal Response](https://mlsanity.com/paper/1407.6981/) (2014)
- [White-box vs Black-box: Bayes Optimal Strategies for Membership Inference](https://mlsanity.com/paper/1908.11229/) (2019)
- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [Learning Continuous Hierarchies in the Lorentz Model of Hyperbolic Geometry](https://mlsanity.com/paper/1806.03417/) (2018)
- [Subsampled R'enyi Differential Privacy and Analytical Moments Accountant](https://mlsanity.com/paper/1808.00087/) (2018)
- [Riemannian Adaptive Optimization Methods](https://mlsanity.com/paper/1810.00760/) (2018)
- [A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups](https://mlsanity.com/paper/2104.09459/) (2021)
- [A Riemannian Network for SPD Matrix Learning](https://mlsanity.com/paper/1608.04233/) (2016)
- [Collecting Telemetry Data Privately](https://mlsanity.com/paper/1712.01524/) (2017)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [Reconstructing Training Data with Informed Adversaries](https://mlsanity.com/paper/2201.04845/) (2022)
- [A micro Lie theory for state estimation in robotics](https://mlsanity.com/paper/1812.01537/) (2018)
- [Manopt, a Matlab toolbox for optimization on manifolds](https://mlsanity.com/paper/1308.5200/) (2013)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Pymanopt: A Python Toolbox for Optimization on Manifolds using Automatic Differentiation](https://mlsanity.com/paper/1603.03236/) (2016)
- [First-order Methods for Geodesically Convex Optimization](https://mlsanity.com/paper/1602.06053/) (2016)
- [geomstats: a Python Package for Riemannian Geometry in Machine Learning](https://mlsanity.com/paper/1805.08308/) (2018)
- [Riemannian SVRG: Fast Stochastic Optimization on Riemannian Manifolds](https://mlsanity.com/paper/1605.07147/) (2016)
- [Building Deep Networks on Grassmann Manifolds](https://mlsanity.com/paper/1611.05742/) (2016)

## Built on by (most cited first, 1 of 1)

- [hyperbolix: Hyperbolic Deep Learning in JAX](https://mlsanity.com/paper/2609.28248/) (2026)

## Abstract (arXiv)

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.