Manifolds

All manifolds share same API. Some manifols may have several implementations of retraction operation, every implementation has a corresponding class.

class geoopt.manifolds.Euclidean(ndim=0)[source]

Simple Euclidean manifold, every coordinate is treated as an independent element.

Parameters:ndim (int) – number of trailing dimensions treated as manifold dimensions. All the operations acting on cuch as inner products, etc will respect the ndim.
component_inner(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\) according to components of the manifold.

The result of the function is same as inner with keepdim=True for all the manifolds except ProductManifold. For this manifold it acts different way computing inner product for each component and then building an output correctly tiling and reshaping the result.

Parameters:
Returns:

inner product component wise (broadcasted)

Return type:

torch.Tensor

Notes

The purpose of this method is better adaptive properties in optimization since ProductManifold will “hide” the structure in public API.

dist(x: torch.Tensor, y: torch.Tensor, *, keepdim=False) → torch.Tensor[source]

Compute distance between 2 points on the manifold that is the shortest path along geodesics.

Parameters:
Returns:

distance between two points

Return type:

torch.Tensor

dist2(x: torch.Tensor, y: torch.Tensor, *, keepdim=False) → torch.Tensor[source]

Compute squared distance between 2 points on the manifold that is the shortest path along geodesics.

Parameters:
Returns:

squared distance between two points

Return type:

torch.Tensor

egrad2rgrad(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Transform gradient computed using autodiff to the correct Riemannian gradient for the point \(x\).

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – gradient to be projected
Returns:

grad vector in the Riemannian manifold

Return type:

torch.Tensor

expmap(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

extra_repr()[source]

Set the extra representation of the module

To print customized extra information, you should reimplement this method in your own modules. Both single-line and multi-line strings are acceptable.

inner(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

logmap(x: torch.Tensor, y: torch.Tensor) → torch.Tensor[source]

Perform an logarithmic map \(\operatorname{Log}_{x}(y)\).

Parameters:
Returns:

tangent vector

Return type:

torch.Tensor

norm(x: torch.Tensor, u: torch.Tensor, *, keepdim=False)[source]

Norm of a tangent vector at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

origin(*size, dtype=None, device=None, seed=42) → geoopt.tensor.ManifoldTensor[source]

Zero point origin.

Parameters:
  • size (shape) – the desired shape
  • device (torch.device) – the desired device
  • dtype (torch.dtype) – the desired dtype
  • seed (int) – ignored
Returns:

Return type:

ManifoldTensor

proju(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

projx(x: torch.Tensor) → torch.Tensor[source]

Project point \(x\) on the manifold.

Parameters:torch.Tensor (x) – point to be projected
Returns:projected point
Return type:torch.Tensor
random(*size, mean=0.0, std=1.0, device=None, dtype=None) → geoopt.tensor.ManifoldTensor

Create a point on the manifold, measure is induced by Normal distribution.

Parameters:
  • size (shape) – the desired shape
  • mean (float|tensor) – mean value for the Normal distribution
  • std (float|tensor) – std value for the Normal distribution
  • device (torch.device) – the desired device
  • dtype (torch.dtype) – the desired dtype
Returns:

random point on the manifold

Return type:

ManifoldTensor

random_normal(*size, mean=0.0, std=1.0, device=None, dtype=None) → geoopt.tensor.ManifoldTensor[source]

Create a point on the manifold, measure is induced by Normal distribution.

Parameters:
  • size (shape) – the desired shape
  • mean (float|tensor) – mean value for the Normal distribution
  • std (float|tensor) – std value for the Normal distribution
  • device (torch.device) – the desired device
  • dtype (torch.dtype) – the desired dtype
Returns:

random point on the manifold

Return type:

ManifoldTensor

retr(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform a retraction from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

transp(x: torch.Tensor, y: torch.Tensor, v: torch.Tensor) → torch.Tensor[source]

Perform vector transport \(\mathfrak{T}_{x\to y}(v)\).

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.Stiefel(**kwargs)[source]

Manifold induced by the following matrix constraint:

\[\begin{split}X^\top X = I\\ X \in \mathrm{R}^{n\times m}\\ n \ge m\end{split}\]
Parameters:canonical (bool) – Use canonical inner product instead of euclidean one (defaults to canonical)
origin(*size, dtype=None, device=None, seed=42) → torch.Tensor[source]

Identity matrix point origin.

Parameters:
  • size (shape) – the desired shape
  • device (torch.device) – the desired device
  • dtype (torch.dtype) – the desired dtype
  • seed (int) – ignored
Returns:

Return type:

ManifoldTensor

projx(x: torch.Tensor) → torch.Tensor[source]

Project point \(x\) on the manifold.

Parameters:torch.Tensor (x) – point to be projected
Returns:projected point
Return type:torch.Tensor
random(*size, dtype=None, device=None) → torch.Tensor

Naive approach to get random matrix on Stiefel manifold.

A helper function to sample a random point on the Stiefel manifold. The measure is non-uniform for this method, but fast to compute.

Parameters:
  • size (shape) – the desired output shape
  • dtype (torch.dtype) – desired dtype
  • device (torch.device) – desired device
Returns:

random point on Stiefel manifold

Return type:

ManifoldTensor

random_naive(*size, dtype=None, device=None) → torch.Tensor[source]

Naive approach to get random matrix on Stiefel manifold.

A helper function to sample a random point on the Stiefel manifold. The measure is non-uniform for this method, but fast to compute.

Parameters:
  • size (shape) – the desired output shape
  • dtype (torch.dtype) – desired dtype
  • device (torch.device) – desired device
Returns:

random point on Stiefel manifold

Return type:

ManifoldTensor

class geoopt.manifolds.CanonicalStiefel(**kwargs)[source]

Stiefel Manifold with Canonical inner product

Manifold induced by the following matrix constraint:

\[\begin{split}X^\top X = I\\ X \in \mathrm{R}^{n\times m}\\ n \ge m\end{split}\]
egrad2rgrad(x: torch.Tensor, u: torch.Tensor) → torch.Tensor

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

expmap(x: torch.Tensor, u: torch.Tensor) → torch.Tensor

Perform a retraction from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

expmap_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Perform a retraction + vector transport at once.

Parameters:
Returns:

transported point and vectors

Return type:

Tuple[torch.Tensor, torch.Tensor]

Notes

Sometimes this is a far more optimal way to preform retraction + vector transport

inner(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

proju(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

retr(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform a retraction from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

retr_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]

Perform a retraction + vector transport at once.

Parameters:
Returns:

transported point and vectors

Return type:

Tuple[torch.Tensor, torch.Tensor]

Notes

Sometimes this is a far more optimal way to preform retraction + vector transport

transp_follow_expmap(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → torch.Tensor

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{retr}(x, u)}(v)\).

This operation is sometimes is much more simpler and can be optimized.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

transp_follow_retr(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → torch.Tensor[source]

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{retr}(x, u)}(v)\).

This operation is sometimes is much more simpler and can be optimized.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.EuclideanStiefel(**kwargs)[source]

Stiefel Manifold with Euclidean inner product

Manifold induced by the following matrix constraint:

\[\begin{split}X^\top X = I\\ X \in \mathrm{R}^{n\times m}\\ n \ge m\end{split}\]
egrad2rgrad(x: torch.Tensor, u: torch.Tensor) → torch.Tensor

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

expmap(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

inner(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

proju(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

retr(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform a retraction from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

transp(x: torch.Tensor, y: torch.Tensor, v: torch.Tensor) → torch.Tensor[source]

Perform vector transport \(\mathfrak{T}_{x\to y}(v)\).

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.EuclideanStiefelExact(**kwargs)[source]

Stiefel Manifold with Euclidean inner product

Manifold induced by the following matrix constraint:

\[\begin{split}X^\top X = I\\ X \in \mathrm{R}^{n\times m}\\ n \ge m\end{split}\]

Notes

The implementation of retraction is an exact exponential map, this retraction will be used in optimization

extra_repr()[source]

Set the extra representation of the module

To print customized extra information, you should reimplement this method in your own modules. Both single-line and multi-line strings are acceptable.

retr(x: torch.Tensor, u: torch.Tensor) → torch.Tensor

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

retr_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Perform an exponential map and vector transport from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

transp_follow_retr(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → torch.Tensor

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{Exp}(x, u)}(v)\).

Here, \(\operatorname{Exp}\) is the best possible approximation of the true exponential map. There are cases when the exact variant is hard or impossible implement, therefore a fallback, non-exact, implementation is used.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.Sphere(intersection: torch.Tensor = None, complement: torch.Tensor = None)[source]

Sphere manifold induced by the following constraint

\[\begin{split}\|x\|=1\\ x \in \mathbb{span}(U)\end{split}\]

where \(U\) can be parametrized with compliment space or intersection.

Parameters:
  • intersection (tensor) – shape (..., dim, K), subspace to intersect with
  • complement (tensor) – shape (..., dim, K), subspace to compliment

See also

SphereExact

dist(x: torch.Tensor, y: torch.Tensor, *, keepdim=False) → torch.Tensor[source]

Compute distance between 2 points on the manifold that is the shortest path along geodesics.

Parameters:
Returns:

distance between two points

Return type:

torch.Tensor

egrad2rgrad(x: torch.Tensor, u: torch.Tensor) → torch.Tensor

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

expmap(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

inner(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

logmap(x: torch.Tensor, y: torch.Tensor) → torch.Tensor[source]

Perform an logarithmic map \(\operatorname{Log}_{x}(y)\).

Parameters:
Returns:

tangent vector

Return type:

torch.Tensor

proju(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

projx(x: torch.Tensor) → torch.Tensor[source]

Project point \(x\) on the manifold.

Parameters:torch.Tensor (x) – point to be projected
Returns:projected point
Return type:torch.Tensor
random(*size, dtype=None, device=None) → torch.Tensor

Uniform random measure on Sphere manifold.

Parameters:
  • size (shape) – the desired output shape
  • dtype (torch.dtype) – desired dtype
  • device (torch.device) – desired device
Returns:

random point on Sphere manifold

Return type:

ManifoldTensor

Notes

In case of projector on the manifold, dtype and device are set automatically and shouldn’t be provided. If you provide them, they are checked to match the projector device and dtype

random_uniform(*size, dtype=None, device=None) → torch.Tensor[source]

Uniform random measure on Sphere manifold.

Parameters:
  • size (shape) – the desired output shape
  • dtype (torch.dtype) – desired dtype
  • device (torch.device) – desired device
Returns:

random point on Sphere manifold

Return type:

ManifoldTensor

Notes

In case of projector on the manifold, dtype and device are set automatically and shouldn’t be provided. If you provide them, they are checked to match the projector device and dtype

retr(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform a retraction from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

transp(x: torch.Tensor, y: torch.Tensor, v: torch.Tensor) → torch.Tensor[source]

Perform vector transport \(\mathfrak{T}_{x\to y}(v)\).

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.SphereExact(intersection: torch.Tensor = None, complement: torch.Tensor = None)[source]

Sphere manifold induced by the following constraint

\[\begin{split}\|x\|=1\\ x \in \mathbb{span}(U)\end{split}\]

where \(U\) can be parametrized with compliment space or intersection.

Parameters:
  • intersection (tensor) – shape (..., dim, K), subspace to intersect with
  • complement (tensor) – shape (..., dim, K), subspace to compliment

See also

Sphere

Notes

The implementation of retraction is an exact exponential map, this retraction will be used in optimization

extra_repr()[source]

Set the extra representation of the module

To print customized extra information, you should reimplement this method in your own modules. Both single-line and multi-line strings are acceptable.

retr(x: torch.Tensor, u: torch.Tensor) → torch.Tensor

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

retr_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Perform an exponential map and vector transport from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

transp_follow_retr(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → torch.Tensor

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{Exp}(x, u)}(v)\).

Here, \(\operatorname{Exp}\) is the best possible approximation of the true exponential map. There are cases when the exact variant is hard or impossible implement, therefore a fallback, non-exact, implementation is used.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.Stereographic(k=0.0, learnable=False)[source]

\(\kappa\)-Stereographic model.

Parameters:k (float|tensor) – sectional curvature \(\kappa\) of the manifold - k<0: Poincaré ball (stereographic projection of hyperboloid) - k>0: Stereographic projection of sphere - k=0: Euclidean geometry

Notes

It is extremely recommended to work with this manifold in double precision.

http://andbloch.github.io/K-Stereographic-Model/ or \kappa-Stereographic Projection model

References

The functions for the mathematics in gyrovector spaces are taken from the following resources:

[1] Ganea, Octavian, Gary Bécigneul, and Thomas Hofmann. “Hyperbolic
neural networks.” Advances in neural information processing systems. 2018.
[2] Bachmann, Gregor, Gary Bécigneul, and Octavian-Eugen Ganea. “Constant
Curvature Graph Convolutional Networks.” arXiv preprint arXiv:1911.05076 (2019).
[3] Skopek, Ondrej, Octavian-Eugen Ganea, and Gary Bécigneul.
“Mixed-curvature Variational Autoencoders.” arXiv preprint arXiv:1911.08411 (2019).
[4] Ungar, Abraham A. Analytic hyperbolic geometry: Mathematical
foundations and applications. World Scientific, 2005.
[5] Albert, Ungar Abraham. Barycentric calculus in Euclidean and
hyperbolic geometry: A comparative introduction. World Scientific, 2010.
dist(x: torch.Tensor, y: torch.Tensor, *, keepdim=False, dim=-1) → torch.Tensor[source]

Compute distance between 2 points on the manifold that is the shortest path along geodesics.

Parameters:
Returns:

distance between two points

Return type:

torch.Tensor

dist2(x: torch.Tensor, y: torch.Tensor, *, keepdim=False, dim=-1) → torch.Tensor[source]

Compute squared distance between 2 points on the manifold that is the shortest path along geodesics.

Parameters:
Returns:

squared distance between two points

Return type:

torch.Tensor

egrad2rgrad(x: torch.Tensor, u: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Transform gradient computed using autodiff to the correct Riemannian gradient for the point \(x\).

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – gradient to be projected
Returns:

grad vector in the Riemannian manifold

Return type:

torch.Tensor

expmap(x: torch.Tensor, u: torch.Tensor, *, project=True, dim=-1) → torch.Tensor[source]

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

expmap_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor, *, dim=-1, project=True) → Tuple[torch.Tensor, torch.Tensor][source]

Perform an exponential map and vector transport from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

inner(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False, dim=-1) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

logmap(x: torch.Tensor, y: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Perform an logarithmic map \(\operatorname{Log}_{x}(y)\).

Parameters:
Returns:

tangent vector

Return type:

torch.Tensor

norm(x: torch.Tensor, u: torch.Tensor, *, keepdim=False, dim=-1) → torch.Tensor[source]

Norm of a tangent vector at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

origin(*size, dtype=None, device=None, seed=42) → geoopt.tensor.ManifoldTensor[source]

Zero point origin.

Parameters:
  • size (shape) – the desired shape
  • device (torch.device) – the desired device
  • dtype (torch.dtype) – the desired dtype
  • seed (int) – ignored
Returns:

random point on the manifold

Return type:

ManifoldTensor

proju(x: torch.Tensor, u: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

projx(x: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Project point \(x\) on the manifold.

Parameters:torch.Tensor (x) – point to be projected
Returns:projected point
Return type:torch.Tensor
random(*size, mean=0, std=1, dtype=None, device=None) → geoopt.tensor.ManifoldTensor

Create a point on the manifold, measure is induced by Normal distribution on the tangent space of zero.

Parameters:
  • size (shape) – the desired shape
  • mean (float|tensor) – mean value for the Normal distribution
  • std (float|tensor) – std value for the Normal distribution
  • dtype (torch.dtype) – target dtype for sample, if not None, should match Manifold dtype
  • device (torch.device) – target device for sample, if not None, should match Manifold device
Returns:

random point on the PoincareBall manifold

Return type:

ManifoldTensor

Notes

The device and dtype will match the device and dtype of the Manifold

random_normal(*size, mean=0, std=1, dtype=None, device=None) → geoopt.tensor.ManifoldTensor[source]

Create a point on the manifold, measure is induced by Normal distribution on the tangent space of zero.

Parameters:
  • size (shape) – the desired shape
  • mean (float|tensor) – mean value for the Normal distribution
  • std (float|tensor) – std value for the Normal distribution
  • dtype (torch.dtype) – target dtype for sample, if not None, should match Manifold dtype
  • device (torch.device) – target device for sample, if not None, should match Manifold device
Returns:

random point on the PoincareBall manifold

Return type:

ManifoldTensor

Notes

The device and dtype will match the device and dtype of the Manifold

retr(x: torch.Tensor, u: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Perform a retraction from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

retr_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor, *, dim=-1) → Tuple[torch.Tensor, torch.Tensor][source]

Perform a retraction + vector transport at once.

Parameters:
Returns:

transported point and vectors

Return type:

Tuple[torch.Tensor, torch.Tensor]

Notes

Sometimes this is a far more optimal way to preform retraction + vector transport

transp(x: torch.Tensor, y: torch.Tensor, v: torch.Tensor, *, dim=-1)[source]

Perform vector transport \(\mathfrak{T}_{x\to y}(v)\).

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

transp_follow_expmap(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor, *, dim=-1, project=True) → torch.Tensor[source]

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{Exp}(x, u)}(v)\).

Here, \(\operatorname{Exp}\) is the best possible approximation of the true exponential map. There are cases when the exact variant is hard or impossible implement, therefore a fallback, non-exact, implementation is used.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

transp_follow_retr(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{retr}(x, u)}(v)\).

This operation is sometimes is much more simpler and can be optimized.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.StereographicExact(k=0.0, learnable=False)[source]

\(\kappa\)-Stereographic model.

Parameters:k (float|tensor) – sectional curvature \(\kappa\) of the manifold - k<0: Poincaré ball (stereographic projection of hyperboloid) - k>0: Stereographic projection of sphere - k=0: Euclidean geometry

Notes

It is extremely recommended to work with this manifold in double precision.

http://andbloch.github.io/K-Stereographic-Model/ or \kappa-Stereographic Projection model

The implementation of retraction is an exact exponential map, this retraction will be used in optimization.

extra_repr()[source]

Set the extra representation of the module

To print customized extra information, you should reimplement this method in your own modules. Both single-line and multi-line strings are acceptable.

retr(x: torch.Tensor, u: torch.Tensor, *, project=True, dim=-1) → torch.Tensor

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

retr_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor, *, dim=-1, project=True) → Tuple[torch.Tensor, torch.Tensor]

Perform an exponential map and vector transport from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

transp_follow_retr(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor, *, dim=-1, project=True) → torch.Tensor

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{Exp}(x, u)}(v)\).

Here, \(\operatorname{Exp}\) is the best possible approximation of the true exponential map. There are cases when the exact variant is hard or impossible implement, therefore a fallback, non-exact, implementation is used.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.PoincareBall(c=1.0, learnable=False)[source]

Poincare ball model.

See more in \kappa-Stereographic Projection model

Parameters:c (float|tensor) – ball’s negative curvature. The parametrization is constrained to have positive c

Notes

It is extremely recommended to work with this manifold in double precision

class geoopt.manifolds.PoincareBallExact(c=1.0, learnable=False)[source]

Poincare ball model.

See more in \kappa-Stereographic Projection model

Parameters:c (float|tensor) – ball’s negative curvature. The parametrization is constrained to have positive c

Notes

It is extremely recommended to work with this manifold in double precision

The implementation of retraction is an exact exponential map, this retraction will be used in optimization.

class geoopt.manifolds.SphereProjection(k=1.0, learnable=False)[source]

Stereographic Projection Spherical model.

See more in \kappa-Stereographic Projection model

Parameters:k (float|tensor) – sphere’s positive curvature. The parametrization is constrained to have positive k

Notes

It is extremely recommended to work with this manifold in double precision

class geoopt.manifolds.SphereProjectionExact(k=1.0, learnable=False)[source]

Stereographic Projection Spherical model.

See more in \kappa-Stereographic Projection model

Parameters:k (float|tensor) – sphere’s positive curvature. The parametrization is constrained to have positive k

Notes

It is extremely recommended to work with this manifold in double precision

The implementation of retraction is an exact exponential map, this retraction will be used in optimization.

class geoopt.manifolds.Scaled(manifold: geoopt.manifolds.base.Manifold, scale=1.0, learnable=False)[source]

Scaled manifold.

Scales all the distances on tha manifold by a constant factor. Scaling may be learnable since the underlying representation is canonical.

Examples

Here is a simple example of radius 2 Sphere

>>> import geoopt, torch, numpy as np
>>> sphere = geoopt.Sphere()
>>> radius_2_sphere = Scaled(sphere, 2)
>>> p1 = torch.tensor([-1., 0.])
>>> p2 = torch.tensor([0., 1.])
>>> np.testing.assert_allclose(sphere.dist(p1, p2), np.pi / 2)
>>> np.testing.assert_allclose(radius_2_sphere.dist(p1, p2), np.pi)
egrad2rgrad(x: torch.Tensor, u: torch.Tensor, **kwargs) → torch.Tensor[source]

Transform gradient computed using autodiff to the correct Riemannian gradient for the point \(x\).

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – gradient to be projected
Returns:

grad vector in the Riemannian manifold

Return type:

torch.Tensor

inner(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False, **kwargs) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

norm(x: torch.Tensor, u: torch.Tensor, *, keepdim=False, **kwargs) → torch.Tensor[source]

Norm of a tangent vector at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

proju(x: torch.Tensor, u: torch.Tensor, **kwargs) → torch.Tensor[source]

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

projx(x: torch.Tensor, **kwargs) → torch.Tensor[source]

Project point \(x\) on the manifold.

Parameters:torch.Tensor (x) – point to be projected
Returns:projected point
Return type:torch.Tensor
random(*size, dtype=None, device=None, **kwargs) → torch.Tensor[source]

Random sampling on the manifold.

The exact implementation depends on manifold and usually does not follow all assumptions about uniform measure, etc.

transp(x: torch.Tensor, y: torch.Tensor, v: torch.Tensor, **kwargs) → torch.Tensor[source]

Perform vector transport \(\mathfrak{T}_{x\to y}(v)\).

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

class geoopt.manifolds.ProductManifold(*manifolds_with_shape)[source]

Product Manifold.

Examples

A Torus

>>> import geoopt
>>> sphere = geoopt.Sphere()
>>> torus = ProductManifold((sphere, 2), (sphere, 2))
component_inner(x: torch.Tensor, u: torch.Tensor, v=None) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\) according to components of the manifold.

The result of the function is same as inner with keepdim=True for all the manifolds except ProductManifold. For this manifold it acts different way computing inner product for each component and then building an output correctly tiling and reshaping the result.

Parameters:
Returns:

inner product component wise (broadcasted)

Return type:

torch.Tensor

Notes

The purpose of this method is better adaptive properties in optimization since ProductManifold will “hide” the structure in public API.

dist(x, y, *, keepdim=False)[source]

Compute distance between 2 points on the manifold that is the shortest path along geodesics.

Parameters:
Returns:

distance between two points

Return type:

torch.Tensor

dist2(x: torch.Tensor, y: torch.Tensor, *, keepdim=False)[source]

Compute squared distance between 2 points on the manifold that is the shortest path along geodesics.

Parameters:
Returns:

squared distance between two points

Return type:

torch.Tensor

egrad2rgrad(x: torch.Tensor, u: torch.Tensor)[source]

Transform gradient computed using autodiff to the correct Riemannian gradient for the point \(x\).

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – gradient to be projected
Returns:

grad vector in the Riemannian manifold

Return type:

torch.Tensor

expmap(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

expmap_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor][source]

Perform an exponential map and vector transport from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

classmethod from_point(*parts, batch_dims=0)[source]

Construct Product manifold from given points.

Parameters:
  • parts (tuple[geoopt.ManifoldTensor]) – Manifold tensors to construct Product manifold from
  • batch_dims (int) – number of first dims to treat as batch dims and not include in the Product manifold
Returns:

Return type:

ProductManifold

inner(x: torch.Tensor, u: torch.Tensor, v=None, *, keepdim=False) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

logmap(x: torch.Tensor, y: torch.Tensor) → torch.Tensor[source]

Perform an logarithmic map \(\operatorname{Log}_{x}(y)\).

Parameters:
Returns:

tangent vector

Return type:

torch.Tensor

origin(*size, dtype=None, device=None, seed=42) → geoopt.tensor.ManifoldTensor[source]

Create some reasonable point on the manifold in a deterministic way.

For some manifolds there may exist e.g. zero vector or some analogy. In case it is possible to define this special point, this point is returned with the desired size. In other case, the returned point is sampled on the manifold in a deterministic way.

Parameters:
  • size (Union[int, Tuple[int]]) – the desired shape
  • device (torch.device) – the desired device
  • dtype (torch.dtype) – the desired dtype
  • seed (Optional[int]) – A parameter controlling deterministic randomness for manifolds that do not provide .origin, but provide .random. (default: 42)
Returns:

Return type:

torch.Tensor

pack_point(*tensors) → torch.Tensor[source]

Construct a tensor representation of a manifold point.

In case of regular manifolds this will return the same tensor. However, for e.g. Product manifold this function will pack all non-batch dimensions.

Parameters:tensors (Tuple[torch.Tensor]) –
Returns:
Return type:torch.Tensor
proju(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

projx(x: torch.Tensor) → torch.Tensor[source]

Project point \(x\) on the manifold.

Parameters:torch.Tensor (x) – point to be projected
Returns:projected point
Return type:torch.Tensor
retr(x: torch.Tensor, u: torch.Tensor) → torch.Tensor[source]

Perform a retraction from point \(x\) with given direction \(u\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

retr_transp(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor)[source]

Perform a retraction + vector transport at once.

Parameters:
Returns:

transported point and vectors

Return type:

Tuple[torch.Tensor, torch.Tensor]

Notes

Sometimes this is a far more optimal way to preform retraction + vector transport

take_submanifold_value(x: torch.Tensor, i: int, reshape=True) → torch.Tensor[source]

Take i’th slice of the ambient tensor and possibly reshape.

Parameters:
  • x (tensor) – Ambient tensor
  • i (int) – submanifold index
  • reshape (bool) – reshape the slice?
Returns:

Return type:

torch.Tensor

transp(x: torch.Tensor, y: torch.Tensor, v: torch.Tensor) → torch.Tensor[source]

Perform vector transport \(\mathfrak{T}_{x\to y}(v)\).

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

transp_follow_expmap(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → torch.Tensor[source]

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{Exp}(x, u)}(v)\).

Here, \(\operatorname{Exp}\) is the best possible approximation of the true exponential map. There are cases when the exact variant is hard or impossible implement, therefore a fallback, non-exact, implementation is used.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

transp_follow_retr(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) → torch.Tensor[source]

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{retr}(x, u)}(v)\).

This operation is sometimes is much more simpler and can be optimized.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

unpack_tensor(tensor: torch.Tensor) → Tuple[torch.Tensor][source]

Construct a point on the manifold.

This method should help to work with product and compound manifolds. Internally all points on the manifold are stored in an intuitive format. However, there might be cases, when this representation is simpler or more efficient to store in a different way that is hard to use in practice.

Parameters:tensor (torch.Tensor) –
Returns:
Return type:torch.Tensor
class geoopt.manifolds.Lorentz(k=1.0, learnable=False)[source]

Lorentz model

Parameters:k (float|tensor) – manifold negative curvature

Notes

It is extremely recommended to work with this manifold in double precision

dist(x: torch.Tensor, y: torch.Tensor, *, keepdim=False, dim=-1) → torch.Tensor[source]

Compute distance between 2 points on the manifold that is the shortest path along geodesics.

Parameters:
Returns:

distance between two points

Return type:

torch.Tensor

egrad2rgrad(x: torch.Tensor, u: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Transform gradient computed using autodiff to the correct Riemannian gradient for the point \(x\).

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – gradient to be projected
Returns:

grad vector in the Riemannian manifold

Return type:

torch.Tensor

expmap(x: torch.Tensor, u: torch.Tensor, *, norm_tan=True, project=True, dim=-1) → torch.Tensor[source]

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

inner(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False, dim=-1) → torch.Tensor[source]

Inner product for tangent vectors at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

logmap(x: torch.Tensor, y: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Perform an logarithmic map \(\operatorname{Log}_{x}(y)\).

Parameters:
Returns:

tangent vector

Return type:

torch.Tensor

norm(u: torch.Tensor, *, keepdim=False, dim=-1) → torch.Tensor[source]

Norm of a tangent vector at point \(x\).

Parameters:
Returns:

inner product (broadcasted)

Return type:

torch.Tensor

origin(*size, dtype=None, device=None, seed=42) → geoopt.tensor.ManifoldTensor[source]

Zero point origin.

Parameters:
  • size (shape) – the desired shape
  • device (torch.device) – the desired device
  • dtype (torch.dtype) – the desired dtype
  • seed (int) – ignored
Returns:

zero point on the manifold

Return type:

ManifoldTensor

proju(x: torch.Tensor, v: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Project vector \(u\) on a tangent space for \(x\), usually is the same as egrad2rgrad().

Parameters:
  • torch.Tensor (u) – point on the manifold
  • torch.Tensor – vector to be projected
Returns:

projected vector

Return type:

torch.Tensor

projx(x: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Project point \(x\) on the manifold.

Parameters:torch.Tensor (x) – point to be projected
Returns:projected point
Return type:torch.Tensor
random_normal(*size, mean=0, std=1, dtype=None, device=None) → geoopt.tensor.ManifoldTensor[source]

Create a point on the manifold, measure is induced by Normal distribution on the tangent space of zero.

Parameters:
  • size (shape) – the desired shape
  • mean (float|tensor) – mean value for the Normal distribution
  • std (float|tensor) – std value for the Normal distribution
  • dtype (torch.dtype) – target dtype for sample, if not None, should match Manifold dtype
  • device (torch.device) – target device for sample, if not None, should match Manifold device
Returns:

random points on Hyperboloid

Return type:

ManifoldTensor

Notes

The device and dtype will match the device and dtype of the Manifold

retr(x: torch.Tensor, u: torch.Tensor, *, norm_tan=True, project=True, dim=-1) → torch.Tensor

Perform an exponential map \(\operatorname{Exp}_x(u)\).

Parameters:
Returns:

transported point

Return type:

torch.Tensor

transp(x: torch.Tensor, y: torch.Tensor, v: torch.Tensor, *, dim=-1) → torch.Tensor[source]

Perform vector transport \(\mathfrak{T}_{x\to y}(v)\).

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor

transp_follow_expmap(x: torch.Tensor, u: torch.Tensor, v: torch.Tensor, *, dim=-1, project=True) → torch.Tensor[source]

Perform vector transport following \(u\): \(\mathfrak{T}_{x\to\operatorname{Exp}(x, u)}(v)\).

Here, \(\operatorname{Exp}\) is the best possible approximation of the true exponential map. There are cases when the exact variant is hard or impossible implement, therefore a fallback, non-exact, implementation is used.

Parameters:
Returns:

transported tensor

Return type:

torch.Tensor