TensorFlow RiemOpt

A library for manifold-constrained optimization in TensorFlow.

Installation

To install a package from PyPI:

pip install tensorflow-riemopt

To add it to a project managed by uv:

uv add tensorflow-riemopt

To install the latest development version from GitHub:

pip install git+https://github.com/master/tensorflow-riemopt.git

To work on the library from a clone of the repository:

git clone https://github.com/master/tensorflow-riemopt.git
cd tensorflow-riemopt
uv sync --group dev

Features

The core package implements concepts in differential geometry, such as manifolds and Riemannian metrics with associated exponential and logarithmic maps, geodesics, retractions, and transports. For manifolds, where closed-form expressions are not available, the library provides numerical approximations.

import tensorflow as tf
import tensorflow_riemopt as riemopt

S = riemopt.manifolds.Sphere()

x = S.projx(tf.constant([0.1, -0.1, 0.1]))
u = S.proju(x, tf.constant([1., 1., 1.]))
v = S.proju(x, tf.constant([-0.7, -1.4, 1.4]))

y = S.exp(x, v)

u_ = S.transp(x, y, u)
v_ = S.transp(x, y, v)

Manifolds

• manifolds.Cholesky - manifold of lower triangular matrices with positive diagonal elements
• manifolds.Euclidean - unconstrained manifold with the Euclidean metric
• manifolds.Grassmannian - manifold of p-dimensional linear subspaces of the n-dimensional space
• manifolds.Hyperboloid - manifold of n-dimensional hyperbolic space embedded in the n+1-dimensional Minkowski space
• manifolds.Poincare - the Poincaré ball model of the hyperbolic space
• manifolds.Product - Cartesian product of manifolds
• manifolds.SPDAffineInvariant - manifold of symmetric positive definite (SPD) matrices endowed with the affine-invariant metric
• manifolds.SPDLogCholesky - SPD manifold with the Log-Cholesky metric
• manifolds.SPDLogEuclidean - SPD manifold with the Log-Euclidean metric
• manifolds.SpecialOrthogonal - manifold of rotation matrices
• manifolds.Sphere - manifold of unit-normalized points
• manifolds.StiefelEuclidean - manifold of orthonormal p-frames in the n-dimensional space endowed with the Euclidean metric
• manifolds.StiefelCanonical - Stiefel manifold with the canonical metric
• manifolds.StiefelCayley - Stiefel manifold the retraction map via an iterative Cayley transform

Optimizers

Constrained optimization algorithms work as drop-in replacements for Keras optimizers for sparse and dense updates in both Eager and Graph modes.

• optimizers.RiemannianSGD - Riemannian Gradient Descent
• optimizers.RiemannianAdam - Riemannian Adam and AMSGrad
• optimizers.ConstrainedRMSprop - Constrained RMSProp

Layers

• layers.ManifoldEmbedding - constrained keras.layers.Embedding layer

Examples

• ANTHEM - Choudhary, Nurendra, Nikhil Rao, Sumeet Katariya, Karthik Subbian, and Chandan K. Reddy. "ANTHEM: Attentive Hyperbolic Entity Model for Product Search." In Proceedings of the Fifteenth ACM International Conference on Web Search and Data Mining, 2022.
• SPDNet - Huang, Zhiwu, and Luc Van Gool. "A Riemannian network for SPD matrix learning." Proceedings of the Thirty-First AAAI Conference on Artificial Intelligence. AAAI Press, 2017.
• LieNet - Huang, Zhiwu, et al. "Deep learning on Lie groups for skeleton-based action recognition." Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.
• GrNet - Huang, Zhiwu, Jiqing Wu, and Luc Van Gool. "Building Deep Networks on Grassmann Manifolds." AAAI. AAAI Press, 2018.

References

If you find TensorFlow RiemOpt useful in your research, please cite:

@misc{smirnov2021tensorflow,
      title={TensorFlow RiemOpt: a library for optimization on Riemannian manifolds},
      author={Oleg Smirnov},
      year={2021},
      eprint={2105.13921},
      archivePrefix={arXiv},
      primaryClass={cs.MS}
}

Contact

For support or development inquiries, contact us at: [email protected]

Contributing

Report bugs and request features in the issue tracker along with a minimal snippet that reproduces the problem. To work on the code, clone the repository and run:

uv sync --locked --group dev
uv run pytest -q tensorflow_riemopt
uv run black --line-length 80 --check tensorflow_riemopt

Acknowledgment

TensorFlow RiemOpt was inspired by many similar projects:

• Manopt, a matlab toolbox for optimization on manifolds
• Pymanopt, a Python toolbox for optimization on manifolds
• Geoopt: Riemannian Optimization in PyTorch
• Geomstats, an open-source Python package for computations and statistics on nonlinear manifolds

License

The code is MIT-licensed.