Geometry API#

Every geometry is a directly importable class. Matrix geometries include the metric in the class name so constructors follow the same convention as GrassmannProjection. Capability metadata distinguishes exact geodesic operations from numerical-local candidates and retraction proxies, and marks whether automatic Hessian conversion is mathematically exact.

Core method contract#

For an array geometry with event shape M.shape, every point or tangent has shape batch_shape + M.shape. Leading batch dimensions may broadcast; event dimensions never broadcast and must match exactly. Product applies the same rule independently to every leaf of its factor pytree.

Member	Inputs	Result and contract
`belongs(x, atol=None)`	ambient point candidate	Boolean over the inferable batch shape. A malformed event shape returns `False`. Open constraints such as positive definiteness and ball membership remain strict.
`project(x)`	ambient array with exactly `M.shape` event axes	A finite manifold point satisfying `belongs(project(x))`. Malformed event axes raise `ValueError`; zero, indefinite, rank-deficient, and noisy values are valid repair inputs.
`is_tangent(x, u, atol=None)`	point and tangent candidate	Boolean over broadcast batch axes. Malformed or incompatible shapes return `False`.
`tangent_project(x, u)`	point and ambient vector	A tangent satisfying `is_tangent(x, tangent_project(x, u))`; malformed event axes raise `ValueError`.
`inner(x, u, v)`	point and two tangent vectors	The Riemannian inner product over broadcast batch axes.
`norm(x, u)`	point and tangent vector	The nonnegative metric norm.
`retr(x, u, t=1)`	point, tangent, and scalar or batched step	A manifold point obtained from the documented retraction.
`exp(x, u)`	point and tangent	Exact exponential, numerical-local operation, or retraction proxy according to `operation_kind("exp")`.
`invretr(x, y)`	base and endpoint	Local inverse of the advertised retraction.
`log(x, y)`	base and endpoint	Exact, numerical-local, or proxy displacement according to `operation_kind("log")`. Genuine cut loci follow each class’s documented branch policy.
`squared_dist(x, y)` / `dist(x, y)`	two points	Scalar per batch element with the status reported by `operation_kind("dist")`.
`transport(x, y, u)`	endpoints and source tangent	A target tangent. `operation_kind("transport")` reports `parallel`, `isometric`, or `vector`.
`egrad_to_rgrad(x, egrad)`	point and ambient gradient	Metric-dual tangent gradient.
`random_point(key, sample_shape=())`	JAX key and static shape	Samples shaped `sample_shape + M.shape`.
`random_tangent(key, x, scale=1, normalize=False)`	JAX key and point	Tangent samples matching `x`; normalization uses the Riemannian norm before `scale`.

Numerical repairs use dtype-aware interior margins. A configured eps is never allowed to disappear through float32 rounding, and fixed-rank repairs place active singular values above the numerical-rank threshold. These repair conventions do not redefine exact maps at genuine cut loci or manifold boundaries.

Capability declarations are conservative. GeometryMixin certifies no exact or isometric operation by default; exact geometries opt into ExactGeometryMixin, while retraction geometries use RetractionGeometryMixin. Product metadata is certified only if every factor provides the corresponding guarantee.

Shared interface and vector geometries#

class geojax.geometry.base.GeometryProtocol(*args, **kwargs)#

Bases: ManifoldProtocol, Protocol

Uniform geometry interface with capability-qualified named operations.

Retraction-only geometries may satisfy this structural protocol through documented compatibility aliases. Use operation_kind to distinguish exact geodesic operations from numerical-local or retraction proxies.

operation_kind(name)#

Return the certified status of a named geometric operation.

Parameters:: name (str)
Return type:: str

exp(x, u)#

Apply the capability-qualified exponential or retraction proxy.

Parameters:

x (Any)
u (Any)

Return type:

Any

log(x, y)#

Apply the capability-qualified logarithm or inverse-retraction proxy.

Parameters:

x (Any)
y (Any)

Return type:

Any

squared_dist(x, y)#

Return the squared capability-qualified point distance.

Parameters:

x (Any)
y (Any)

Return type:

Any

dist(x, y)#

Return the capability-qualified point distance.

Parameters:

x (Any)
y (Any)

Return type:

Any

pair_mean(x, y)#

Return the midpoint-like construction supported by the geometry.

Parameters:

x (Any)
y (Any)

Return type:

Any

class geojax.geometry.base.ManifoldProtocol(*args, **kwargs)#

Bases: Protocol

Retraction-based interface consumed by manifold optimizers.

This is the GeoJAX counterpart of Manopt’s core manifold structure. Exact geodesic maps are deliberately not required.

belongs(x, atol=None)#

Return one membership boolean per broadcast batch element.

Parameters:

x (Any)
atol (float | None)

Return type:

Any

project(x)#

Repair an ambient array with event shape shape into the manifold.

Parameters:: x (Any)
Return type:: Any

is_tangent(x, u, atol=None)#

Return whether u satisfies the tangent constraints at x.

Parameters:

x (Any)
u (Any)
atol (float | None)

Return type:

Any

tangent_project(x, u)#

Project an ambient vector onto the represented tangent space at x.

Parameters:

x (Any)
u (Any)

Return type:

Any

inner(x, u, v)#

Evaluate the Riemannian inner product at x.

Parameters:

x (Any)
u (Any)
v (Any)

Return type:

Any

norm(x, u)#

Evaluate the Riemannian norm induced by inner.

Parameters:

x (Any)
u (Any)

Return type:

Any

lincomb(x, *terms)#

Form coefficient/vector pairs in the tangent space at x.

Parameters:

x (Any)
terms (Any)

Return type:

Any

retr(x, u, t=1.0)#

Retract t * u from x to a manifold point.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

invretr(x, y)#

Return the documented local inverse-retraction displacement.

Parameters:

x (Any)
y (Any)

Return type:

Any

transport(x, y, u)#

Move u from T_x M to T_y M using the advertised transport.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

egrad_to_rgrad(x, egrad)#

Convert an ambient Euclidean gradient to the metric-dual tangent gradient.

Parameters:

x (Any)
egrad (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Convert an ambient Hessian-vector product using the advertised capability.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

random_point(key, sample_shape=())#

Sample points with shape sample_shape + shape.

Parameters:

key (Any)
sample_shape (int | Sequence[int] | Tuple[int, ...])

Return type:

Any

random_tangent(key, x, *, scale=1.0, normalize=False)#

Sample tangent vectors, optionally normalizing before applying scale.

Parameters:

key (Any)
x (Any)
scale (float | Any)
normalize (bool)

Return type:

Any

class geojax.geometry.base.GeometryMixin#

Bases: object

Default helpers for geometries with the GeoJAX protocol.

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Geometries with a more direct or more stable formula should override this method. The logarithm-based default is also meaningful for retraction geometries, where it inherits the documented proxy semantics of log and dist.

Parameters:

x (Any)
y (Any)

Return type:

Any

retr(x, u, t=1.0)#

Default retraction: use the exponential map.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

invretr(x, y)#

Default inverse retraction: use the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

class geojax.geometry.base.ExactGeometryMixin#

Bases: GeometryMixin

Opt-in defaults for geometries with certified exact geodesic operations.

class geojax.geometry.base.RetractionGeometryMixin#

Bases: GeometryMixin

Compatibility maps for manifolds known only through a retraction.

Subclasses implement retr and invretr. exp, log and dist remain available for algorithms that need point differences, but their metadata identifies them as proxies rather than genuine geodesic maps.

exp(x, u)#

Retraction proxy for the exponential map.

Parameters:

x (Any)
u (Any)

Return type:

Any

log(x, y)#

Inverse-retraction proxy for the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

dist(x, y)#

Local inverse-retraction norm; not a geodesic distance.

Parameters:

x (Any)
y (Any)

Return type:

Any

transport(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

transp(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

class geojax.geometry.Euclidean(size, *, atol=1e-06)#

Bases: ExactGeometryMixin

Flat Euclidean geometry.

Parameters:

size (tuple[int, ...]) – Shape of one unbatched point. size=5 is interpreted as (5,).
atol (float) – Tolerance used in shape/tangency checks.

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

retr(x, u, t=1.0)#

Default retraction: use the exponential map.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

class geojax.geometry.Oblique(size, *, atol=1e-06, eps=1e-12)#

Bases: ExactGeometryMixin

Matrices whose columns have unit Euclidean norm.

Oblique(size=(n, m)) is the efficient matrix representation of a product of m copies of the sphere S^(n-1).

Parameters:

size (tuple[int, int])
atol (float)
eps (float)

squared_dist(X, Y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

X (Any)
Y (Any)

Return type:

Any

class geojax.geometry.ProbabilitySimplex(size, *, atol=1e-06, eps=1e-10)#

Bases: ExactGeometryMixin

Interior probability simplex with the Fisher–Rao metric.

Parameters:

size (int)
atol (float)
eps (float)

retr(p, u, t=1.0)#

Positive normalized-addition retraction used by optimizers.

Parameters:

p (Any)
u (Any)
t (float | Any)

Return type:

Any

squared_dist(p, q)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

p (Any)
q (Any)

Return type:

Any

class geojax.geometry.Sphere(size, *, atol=1e-06, eps=1e-12)#

Bases: ExactGeometryMixin

Canonical geometry of the unit sphere S^d in R^{d+1}.

Parameters:

size (int) – Ambient Euclidean dimension, equal to d + 1 for S^d.
atol (float) – Absolute tolerance used by membership and tangency checks.
eps (float) – Small positive number used in numerically stable divisions.

property dim: int#: Intrinsic dimension d of S^d.

property shape: tuple[int]#: Shape of one unbatched point.

belongs(x, atol=None)#

Check whether x lies on the unit sphere.

Parameters:

x (Any)
atol (float | None)

Return type:

Any

is_tangent(x, u, atol=None)#

Check whether u is tangent at x, i.e. <x, u> = 0.

Parameters:

x (Any)
u (Any)
atol (float | None)

Return type:

Any

project(x)#

Normalize a nonzero ambient vector to the unit sphere.

This is for initialization or numerical repair. The core geometry methods use exact sphere formulas.

Parameters:: x (Any)
Return type:: Any

normalize(x)#

Normalize a nonzero ambient vector to the unit sphere.

This is for initialization or numerical repair. The core geometry methods use exact sphere formulas.

Parameters:: x (Any)
Return type:: Any

tangent_project(x, u)#

Orthogonally project an ambient vector u to T_x S^d.

Formula: Proj_x(u) = u - <x, u> x.

Parameters:

x (Any)
u (Any)

Return type:

Any

projection(x, u)#

Orthogonally project an ambient vector u to T_x S^d.

Formula: Proj_x(u) = u - <x, u> x.

Parameters:

x (Any)
u (Any)

Return type:

Any

proj(x, u)#

Orthogonally project an ambient vector u to T_x S^d.

Formula: Proj_x(u) = u - <x, u> x.

Parameters:

x (Any)
u (Any)

Return type:

Any

to_tangent(x, u)#

Orthogonally project an ambient vector u to T_x S^d.

Formula: Proj_x(u) = u - <x, u> x.

Parameters:

x (Any)
u (Any)

Return type:

Any

inner(x, u, v)#

Canonical Riemannian inner product on T_x S^d.

Parameters:

x (Any)
u (Any)
v (Any)

Return type:

Any

norm(x, u)#

Canonical Riemannian norm of a tangent vector.

Parameters:

x (Any)
u (Any)

Return type:

Any

exp(x, u)#

Riemannian exponential map on the unit sphere.

For u in T_x S^d and r = ||u||,: Exp_x(u) = cos(r) x + sin(r) u / r.

Parameters:

x (Any)
u (Any)

Return type:

Any

log(x, y)#

Riemannian logarithm map on the unit sphere.

For y not antipodal to x, the map is

Log_x(y) = theta / sin(theta) * (y - cos(theta) x),

where theta = arccos(<x, y>).

At the antipode the logarithm is not unique; this implementation returns NaN there.

Parameters:

x (Any)
y (Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance with a finite coincident-point gradient.

Parameters:

x (Any)
y (Any)

Return type:

Any

dist(x, y)#

Geodesic distance on the unit sphere.

Parameters:

x (Any)
y (Any)

Return type:

Any

transport(x, y, u)#

Parallel transport along the unique shortest geodesic from x to y.

For x != -y,: PT_{x->y}(u) = u - (<u, y> / (1 + <x, y>)) (x + y).

At antipodal endpoints the shortest geodesic is not unique; this implementation returns NaN there.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

geodesic_flow(x, v, t=1.0)#

Exact kinetic/geodesic flow on the sphere.

Given x in S^d and v in T_x S^d, returns (x_t, v_t), where x_t is the geodesic position at time t and v_t is its velocity at time t.

This is the key primitive for Geodesic Monte Carlo.

Parameters:

x (Any)
v (Any)
t (float | Any)

Return type:

tuple[Any, Any]

egrad_to_rgrad(x, egrad)#

Convert an ambient Euclidean gradient to a Riemannian gradient.

Parameters:

x (Any)
egrad (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Convert an ambient Hessian product using the sphere shape operator.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

random_point(key, sample_shape=())#

Sample uniformly from the unit sphere using normalized Gaussians.

Parameters:

key (Any)
sample_shape (int | Sequence[int] | Tuple[int, ...])

Return type:

Any

random_tangent(key, x, scale=1.0, normalize=False)#

Sample a tangent Gaussian at x.

By default this samples Proj_x(z), z ~ N(0, I), which is the canonical Gaussian on T_x S^d under the round metric. If normalize=True, the resulting tangent vector is normalized before scaling.

Parameters:

key (Any)
x (Any)
scale (float | Any)
normalize (bool)

Return type:

Any

class geojax.geometry.SphereExtrinsic(size, *, atol=1e-06, eps=1e-12)#

Bases: Sphere

Sphere geometry equipped with its identity equivariant embedding.

The embedding j(x) = x maps the unit sphere into its ambient Euclidean space. Since its differential is also the identity, the pullback metric is the round metric already implemented by Sphere; the Riemannian methods inner, exp, log, dist, and transport therefore remain unchanged.

This class adds the genuinely extrinsic operations associated with the embedding: Euclidean chordal distance, nearest-point projection from the embedding space, and the projected ambient mean. The extrinsic mean is undefined when the weighted ambient mean is zero; extrinsic_mean returns NaNs in that case.

Parameters:

size (int)
atol (float)
eps (float)

embed(x)#

Apply the identity embedding j(x) = x.

Parameters:: x (Any)
Return type:: Any

to_embedding(x)#

Apply the identity embedding j(x) = x.

Parameters:: x (Any)
Return type:: Any

from_embedding(z)#

Project an ambient vector to its nearest point on the sphere.

Parameters:: z (Any)
Return type:: Any

project_embedding(z)#

Project an ambient vector to its nearest point on the sphere.

Parameters:: z (Any)
Return type:: Any

squared_chordal_dist(x, y)#

Squared Euclidean distance between the embedded points.

Parameters:

x (Any)
y (Any)

Return type:

Any

chordal_dist(x, y)#

Euclidean chordal distance ||j(x) - j(y)||.

Parameters:

x (Any)
y (Any)

Return type:

Any

embedding_dist(x, y)#

Euclidean chordal distance ||j(x) - j(y)||.

Parameters:

x (Any)
y (Any)

Return type:

Any

extrinsic_mean(points, weights=None)#

Project the weighted ambient mean back to the sphere.

Parameters:

points (Any)
weights (Any | None)

Return type:

Any

class geojax.geometry.PoincareBall(size, *, atol=1e-06, eps=1e-10)#

Bases: ExactGeometryMixin

Poincaré ball model of curvature-minus-one hyperbolic space.

Parameters:

size (int)
atol (float)
eps (float)

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

class geojax.geometry.Grassmann(size, *, atol=1e-06, eps=1e-12)#

Bases: ExactGeometryMixin

Canonical Grassmann geometry Gr(rank, size).

Parameters:

size (tuple[int, int]) – Pair (ambient_dim, rank).
atol (float)
eps (float)

lincomb(X, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

X (Any)
terms (Any)

Return type:

Any

exp(X, U)#

Grassmann exponential map in ONB coordinates.

For a horizontal tangent U, the skew generator

Omega = U @ X.T - X @ U.T

produces the exact geodesic representative expm(Omega) @ X. This formulation is equivalent to the principal-angle SVD formula but has a well-defined JVP when U is zero or has repeated singular values.

Parameters:

X (Any)
U (Any)

Return type:

Any

retr(X, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

X (Any)
U (Any)
t (float | Any)

Return type:

Any

log(X, Y)#

Grassmann logarithm map.

This is well defined away from the cut locus; numerically this requires X^T Y to be nonsingular, i.e. no principal angle equal to pi/2.

Parameters:

X (Any)
Y (Any)

Return type:

Any

squared_dist(X, Y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

X (Any)
Y (Any)

Return type:

Any

transport(X, Y, Z)#

Parallel transport along the shortest geodesic from X to Y.

This uses the standard Grassmann formula based on the SVD of eta = Log_X(Y). The returned vector is represented at the supplied endpoint representative Y and is horizontally projected there.

Parameters:

X (Any)
Y (Any)
Z (Any)

Return type:

Any

transp(X, Y, Z)#

Parallel transport along the shortest geodesic from X to Y.

This uses the standard Grassmann formula based on the SVD of eta = Log_X(Y). The returned vector is represented at the supplied endpoint representative Y and is horizontally projected there.

Parameters:

X (Any)
Y (Any)
Z (Any)

Return type:

Any

projector(X)#

Return the rank-rank orthogonal projector XX^T.

Parameters:: X (Any)
Return type:: Any

class geojax.geometry.GrassmannProjection(size, *, atol=1e-06, eps=1e-12)#

Bases: ExactGeometryMixin

Projection-embedded Grassmann geometry.

Public points are (n, k) orthonormal frames X, exactly as for Grassmann. Operations are evaluated after the equivariant embedding

j([X]) = X @ X.T.

The embedded point is a symmetric rank-k projector P. A horizontal frame tangent U is embedded as

dj_X(U) = U @ X.T + X @ U.T.

The Riemannian metric is the normalized Frobenius metric

inner(P, H, K) = 0.5 * trace(H @ K),

which makes j an isometric embedding. Public tangent vectors are returned in (n, k) horizontal form. Internally, exponential maps and transport use projector tangents and matrix exponentials.

Returning from an embedded matrix uses the top k eigenvectors of its symmetric part. When a reference frame is supplied, an orthogonal Procrustes alignment fixes the otherwise arbitrary right-orthogonal gauge. Use chordal_dist() for the extrinsic projection distance ||j(X) - j(Y)||_F / sqrt(2).

Parameters:

size (tuple[int, int]) – Pair (ambient_dim, rank). Public points have the same shape.
atol (float)
eps (float)

project(A)#

Project an ambient (n, k) matrix to an orthonormal frame.

Parameters:: A (Any)
Return type:: Any

normalize(A)#

Project an ambient (n, k) matrix to an orthonormal frame.

Parameters:: A (Any)
Return type:: Any

tangent_project(X, A)#

Project an ambient frame vector through projector coordinates.

Parameters:

X (Any)
A (Any)

Return type:

Any

projection(X, A)#

Project an ambient frame vector through projector coordinates.

Parameters:

X (Any)
A (Any)

Return type:

Any

proj(X, A)#

Project an ambient frame vector through projector coordinates.

Parameters:

X (Any)
A (Any)

Return type:

Any

to_tangent(X, A)#

Project an ambient frame vector through projector coordinates.

Parameters:

X (Any)
A (Any)

Return type:

Any

embed(X)#

Map an orthonormal frame X to its projector X @ X.T.

Parameters:: X (Any)
Return type:: Any

from_frame(X)#

Map an orthonormal frame X to its projector X @ X.T.

Parameters:: X (Any)
Return type:: Any

to_frame(P, reference=None)#

Recover a frame using a top-eigenspace decomposition.

If reference is supplied, the recovered frame is right-aligned to it by the orthogonal Procrustes solution.

Parameters:

P (Any)
reference (Any | None)

Return type:

Any

from_embedding(P, reference=None)#

Recover a frame using a top-eigenspace decomposition.

If reference is supplied, the recovered frame is right-aligned to it by the orthogonal Procrustes solution.

Parameters:

P (Any)
reference (Any | None)

Return type:

Any

project_embedding(A, reference=None)#

Project a symmetric matrix to a frame for its nearest rank-k projector.

Parameters:

A (Any)
reference (Any | None)

Return type:

Any

embed_tangent(X, U)#

Map a horizontal frame tangent to a symmetric projector tangent.

Parameters:

X (Any)
U (Any)

Return type:

Any

from_projector_tangent(X, H)#

Map a projector tangent back to the horizontal frame gauge at X.

Parameters:

X (Any)
H (Any)

Return type:

Any

exp(X, U)#

Evaluate the Riemannian exponential in projector coordinates.

Parameters:

X (Any)
U (Any)

Return type:

Any

retr(X, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

X (Any)
U (Any)
t (float | Any)

Return type:

Any

log(X, Y)#

Riemannian logarithm, defined away from the Grassmann cut locus.

Parameters:

X (Any)
Y (Any)

Return type:

Any

squared_dist(X, Y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

X (Any)
Y (Any)

Return type:

Any

dist(X, Y)#

Intrinsic principal-angle geodesic distance.

Parameters:

X (Any)
Y (Any)

Return type:

Any

geodesic_dist(X, Y)#

Intrinsic principal-angle geodesic distance.

Parameters:

X (Any)
Y (Any)

Return type:

Any

squared_chordal_dist(X, Y)#

Squared extrinsic distance 0.5 * ||j(X) - j(Y)||_F^2.

Parameters:

X (Any)
Y (Any)

Return type:

Any

chordal_dist(X, Y)#

Extrinsic projection distance ||j(X) - j(Y)||_F / sqrt(2).

Parameters:

X (Any)
Y (Any)

Return type:

Any

projection_dist(X, Y)#

Extrinsic projection distance ||j(X) - j(Y)||_F / sqrt(2).

Parameters:

X (Any)
Y (Any)

Return type:

Any

transport(X, Y, U)#

Parallel transport evaluated in projector coordinates.

Parameters:

X (Any)
Y (Any)
U (Any)

Return type:

Any

transp(X, Y, U)#

Parallel transport evaluated in projector coordinates.

Parameters:

X (Any)
Y (Any)
U (Any)

Return type:

Any

ehess_to_rhess(X, egrad, ehess_vec, U)#

Convert an ambient frame Hessian-vector product to Grassmann form.

Parameters:

X (Any)
egrad (Any)
ehess_vec (Any)
U (Any)

Return type:

Any

extrinsic_mean(points, weights=None)#

Project the weighted mean projector back to an orthonormal frame.

Parameters:

points (Any)
weights (Any | None)

Return type:

Any

class geojax.geometry.Stiefel(size, *, atol=1e-06, eps=1e-12, log_maxiter=32, log_tol=1e-09, log_damping=1e-06)#

Bases: _StiefelBase

Stiefel manifold with the canonical quotient metric.

A point is an n x k orthonormal frame. The metric is the quotient metric induced by O(n) -> O(n) / O(n-k):

g_X(U, V) = trace(U.T @ (I - 0.5 * X @ X.T) @ V).

The exponential map is exact. The logarithm is computed by differentiable damped endpoint shooting; use log_with_info() to inspect convergence. Its capability status is numerical-local because endpoint convergence does not certify a globally shortest geodesic.

Parameters:

size (tuple[int, int])
atol (float)
eps (float)
log_maxiter (int)
log_tol (float)
log_damping (float)

exp(X, U)#

Evaluate the exact canonical-metric exponential map.

Parameters:

X (Any)
U (Any)

Return type:

Any

egrad_to_rgrad(X, egrad)#

Convert an ambient Euclidean gradient for the canonical metric.

Parameters:

X (Any)
egrad (Any)

Return type:

Any

egrad2rgrad(X, egrad)#

Convert an ambient Euclidean gradient for the canonical metric.

Parameters:

X (Any)
egrad (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

invretr(x, y)#

Default inverse retraction: use the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

property k: int#: Number of orthonormal frame columns.

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log(X, Y)#

Return the selected local logarithm, or nonfinite values on failure.

Parameters:

X (Any)
Y (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

log_with_info(X, Y)#

Return a numerical logarithm and endpoint-shooting diagnostics.

The tangent result is the best iterate even if the solver does not converge. Check info.converged before using it as a logarithm.

Parameters:

X (Any)
Y (Any)

Return type:

tuple[Any, StiefelLogInfo]

property n: int#: Ambient dimension.

normalize(A)#

Return the nearest orthonormal frame in Frobenius norm.

Parameters:: A (Any)
Return type:: Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

proj(X, A)#

Orthogonally project an ambient matrix onto the frame tangent space.

Parameters:

X (Any)
A (Any)

Return type:

Any

project(A)#

Return the nearest orthonormal frame in Frobenius norm.

Parameters:: A (Any)
Return type:: Any

projection(X, A)#

Orthogonally project an ambient matrix onto the frame tangent space.

Parameters:

X (Any)
A (Any)

Return type:

Any

retr(x, u, t=1.0)#

Default retraction: use the exponential map.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

squared_dist(X, Y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

X (Any)
Y (Any)

Return type:

Any

tangent_project(X, A)#

Orthogonally project an ambient matrix onto the frame tangent space.

Parameters:

X (Any)
A (Any)

Return type:

Any

to_tangent(X, A)#

Orthogonally project an ambient matrix onto the frame tangent space.

Parameters:

X (Any)
A (Any)

Return type:

Any

transp(X, Y, U)#

Apply an isometric group-action vector transport from X to Y.

This transport is tangent and exactly metric-preserving, but it is not the Levi-Civita parallel transport for either Stiefel metric.

Parameters:

X (Any)
Y (Any)
U (Any)

Return type:

Any

transport(X, Y, U)#

Apply an isometric group-action vector transport from X to Y.

This transport is tangent and exactly metric-preserving, but it is not the Levi-Civita parallel transport for either Stiefel metric.

Parameters:

X (Any)
Y (Any)
U (Any)

Return type:

Any

class geojax.geometry.StiefelEuclidean(size, *, atol=1e-06, eps=1e-12, log_maxiter=32, log_tol=1e-09, log_damping=1e-06)#

Bases: _StiefelBase

Stiefel manifold with the metric induced by its Euclidean embedding.

Points and tangents use the same n x k frame representation as Stiefel, but g_X(U, V) = trace(U.T @ V). Its geodesics therefore differ whenever the tangent has a component X @ A with A skew.

Parameters:

size (tuple[int, int])
atol (float)
eps (float)
log_maxiter (int)
log_tol (float)
log_damping (float)

exp(X, U)#

Evaluate the exact embedded-Euclidean exponential map.

Parameters:

X (Any)
U (Any)

Return type:

Any

ehess_to_rhess(X, egrad, ehess_vec, U)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

X (Any)
egrad (Any)
ehess_vec (Any)
U (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

invretr(x, y)#

Default inverse retraction: use the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

property k: int#: Number of orthonormal frame columns.

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log(X, Y)#

Return the selected local logarithm, or nonfinite values on failure.

Parameters:

X (Any)
Y (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

log_with_info(X, Y)#

Return a numerical logarithm and endpoint-shooting diagnostics.

The tangent result is the best iterate even if the solver does not converge. Check info.converged before using it as a logarithm.

Parameters:

X (Any)
Y (Any)

Return type:

tuple[Any, StiefelLogInfo]

property n: int#: Ambient dimension.

normalize(A)#

Return the nearest orthonormal frame in Frobenius norm.

Parameters:: A (Any)
Return type:: Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

proj(X, A)#

Orthogonally project an ambient matrix onto the frame tangent space.

Parameters:

X (Any)
A (Any)

Return type:

Any

project(A)#

Return the nearest orthonormal frame in Frobenius norm.

Parameters:: A (Any)
Return type:: Any

projection(X, A)#

Orthogonally project an ambient matrix onto the frame tangent space.

Parameters:

X (Any)
A (Any)

Return type:

Any

retr(x, u, t=1.0)#

Default retraction: use the exponential map.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

squared_dist(X, Y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

X (Any)
Y (Any)

Return type:

Any

tangent_project(X, A)#

Orthogonally project an ambient matrix onto the frame tangent space.

Parameters:

X (Any)
A (Any)

Return type:

Any

to_tangent(X, A)#

Orthogonally project an ambient matrix onto the frame tangent space.

Parameters:

X (Any)
A (Any)

Return type:

Any

transp(X, Y, U)#

Apply an isometric group-action vector transport from X to Y.

This transport is tangent and exactly metric-preserving, but it is not the Levi-Civita parallel transport for either Stiefel metric.

Parameters:

X (Any)
Y (Any)
U (Any)

Return type:

Any

transport(X, Y, U)#

Apply an isometric group-action vector transport from X to Y.

This transport is tangent and exactly metric-preserving, but it is not the Levi-Civita parallel transport for either Stiefel metric.

Parameters:

X (Any)
Y (Any)
U (Any)

Return type:

Any

class geojax.geometry.StiefelLogInfo(converged, iterations, residual_norm, step_norm)#

Bases: NamedTuple

Convergence data returned by Stiefel.log_with_info.

log returns nonfinite values when converged is false. Use log_with_info to inspect the best shooting iterate and these diagnostics. Convergence certifies endpoint agreement, not global shortestness.

Parameters:

converged (Any)
iterations (Any)
residual_norm (Any)
step_norm (Any)

converged: Any#: Alias for field number 0

iterations: Any#: Alias for field number 1

residual_norm: Any#: Alias for field number 2

step_norm: Any#: Alias for field number 3

class geojax.geometry.GeneralizedStiefel(size, *, metric, atol=1e-06, eps=1e-10, log_maxiter=32, log_tol=1e-09)#

Bases: ExactGeometryMixin

Frames satisfying X.T @ metric @ X = I.

The metric trace(U.T @ metric @ V) is the pullback of the embedded Euclidean Stiefel metric under X -> metric^(1/2) X. Its exact exponential and numerical-local logarithm are pulled back through this isometry.

Parameters:

size (tuple[int, int])
metric (Any)
atol (float)
eps (float)
log_maxiter (int)
log_tol (float)

log_with_info(X, Y)#

Return the pulled-back local logarithm candidate and shooting diagnostics.

Parameters:

X (Any)
Y (Any)

Return type:

tuple[Any, StiefelLogInfo]

squared_dist(X, Y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

X (Any)
Y (Any)

Return type:

Any

ehess_to_rhess(X, egrad, ehess_vec, U)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

X (Any)
egrad (Any)
ehess_vec (Any)
U (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

invretr(x, y)#

Default inverse retraction: use the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

retr(x, u, t=1.0)#

Default retraction: use the exponential map.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

class geojax.geometry.GeneralizedGrassmann(size, *, metric, atol=1e-06, eps=1e-10)#

Bases: ExactGeometryMixin

Generalized Grassmann geometry for metric-orthonormal subspaces.

Parameters:

size (tuple[int, int])
metric (Any)
atol (float)
eps (float)

squared_dist(X, Y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

X (Any)
Y (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

invretr(x, y)#

Default inverse retraction: use the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

retr(x, u, t=1.0)#

Default retraction: use the exponential map.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

class geojax.geometry.Hyperboloid(size, *, atol=1e-06, eps=1e-12)#

Bases: ExactGeometryMixin

Upper-sheet hyperboloid in ambient Minkowski space R^size.

Parameters:

size (int)
atol (float)
eps (float)

squared_dist(x, y)#

Squared hyperbolic distance with a finite derivative at coincidence.

Parameters:

x (Any)
y (Any)

Return type:

Any

to_poincare(x)#

Convert hyperboloid points to the Poincare ball model.

Parameters:: x (Any)
Return type:: Any

from_poincare(point)#

Convert points in the open Poincare ball to the hyperboloid.

Parameters:: point (Any)
Return type:: Any

class geojax.geometry.Torus(size, *, atol=1e-06)#

Bases: ExactGeometryMixin

Flat d-torus represented by angles in [-pi, pi).

Parameters:

size (int)
atol (float)

retr(x, u, t=1.0)#

Default retraction: use the exponential map.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

class geojax.geometry.Product(factors)#

Bases: GeometryMixin

Direct product of manifold geometries.

factors may be any JAX pytree whose leaves are geometry objects. Product points and tangent vectors must have the same tree structure.

Parameters:: factors (Any)

property exp_is_exact: bool#: Returns True when the argument is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property log_is_exact: bool#: Returns True when the argument is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property dist_is_exact: bool#: Returns True when the argument is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property transport_is_isometric: bool#: Returns True when the argument is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property transport_is_parallel: bool#: Returns True when the argument is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property hessian_conversion_is_exact: bool#: Returns True when the argument is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property riemannian_gradient_jvp_is_exact: bool#: Returns True when the argument is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

retr(x, u, t=1.0)#

Default retraction: use the exponential map.

Parameters:

x (Any)
u (Any)
t (float | Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

Matrix Lie groups#

class geojax.geometry.SpecialOrthogonal(size, *, atol=1e-06, eps=1e-12)#

Bases: ExactGeometryMixin

Rotation group SO(n) with the Frobenius bi-invariant metric.

A point is an n x n matrix R satisfying R.T @ R = I and det(R) = 1. A tangent vector at R is represented in ambient form as R @ Omega for a skew-symmetric matrix Omega.

Parameters:

size (int)
atol (float)
eps (float)

retr(R, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

R (Any)
U (Any)
t (float | Any)

Return type:

Any

squared_dist(R, Q)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

R (Any)
Q (Any)

Return type:

Any

transport(R, Q, U)#

Parallel transport along the selected shortest geodesic.

Parameters:

R (Any)
Q (Any)
U (Any)

Return type:

Any

transp(R, Q, U)#

Parallel transport along the selected shortest geodesic.

Parameters:

R (Any)
Q (Any)
U (Any)

Return type:

Any

ehess_to_rhess(R, egrad, ehess_vec, U)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

R (Any)
egrad (Any)
ehess_vec (Any)
U (Any)

Return type:

Any

group_exp(omega)#

Lie-group exponential from a skew matrix at the identity.

Parameters:: omega (Any)
Return type:: Any

group_log(R)#

Principal Lie-group logarithm, undefined at rotations by pi.

Parameters:: R (Any)
Return type:: Any

class geojax.geometry.SpecialEuclidean(size, *, atol=1e-06, eps=1e-12)#

Bases: ExactGeometryMixin

Rigid-motion group SE(n) with its canonical product metric.

Points use homogeneous matrices [[R, t], [0, 1]]. The metric is the direct product of the Frobenius bi-invariant metric on SO(n) and the Euclidean metric on translations. Its Riemannian exponential is distinct from the Lie-group exponential except for special tangent directions.

Parameters:

size (int)
atol (float)
eps (float)

retr(G, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

G (Any)
U (Any)
t (float | Any)

Return type:

Any

squared_dist(G, H)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

G (Any)
H (Any)

Return type:

Any

ehess_to_rhess(G, egrad, ehess_vec, U)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

G (Any)
egrad (Any)
ehess_vec (Any)
U (Any)

Return type:

Any

group_exp(tangent_at_identity)#

Lie-group exponential of a homogeneous Lie-algebra matrix.

Parameters:: tangent_at_identity (Any)
Return type:: Any

group_log(G)#

Principal Lie-group logarithm, undefined at rotations by pi.

Parameters:: G (Any)
Return type:: Any

Positive-definite matrices#

class geojax.geometry.SPDLogEuclidean(size, *, atol=1e-06, eps=1e-10)#

Bases: ExactGeometryMixin

Log-Euclidean geometry on SPD(size).

The logarithm map log: SPD(n) -> Sym(n) is an isometry. Therefore distances and geodesics are Euclidean after applying the matrix logarithm. Points are represented as SPD matrices of shape (size, size). Tangent vectors are symmetric matrices of the same shape.

Parameters:

size (tuple[int, int])
atol (float)
eps (float)

retr(P, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

P (Any)
U (Any)
t (float | Any)

Return type:

Any

squared_dist(P, Q)#

Squared Euclidean distance between matrix-log coordinates.

Parameters:

P (Any)
Q (Any)

Return type:

Any

lincomb(P, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

P (Any)
terms (Any)

Return type:

Any

class geojax.geometry.SPDAffineInvariant(size, *, atol=1e-06, eps=1e-10)#

Bases: ExactGeometryMixin

Affine-invariant geometry on SPD(size).

The metric is

g_P(U,V) = tr(P^{-1} U P^{-1} V).

This is the canonical symmetric-space geometry of GL(n)/O(n).

Parameters:

size (tuple[int, int])
atol (float)
eps (float)

retr(P, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

P (Any)
U (Any)
t (float | Any)

Return type:

Any

squared_dist(P, Q)#

Squared affine-invariant distance.

Parameters:

P (Any)
Q (Any)

Return type:

Any

lincomb(P, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

P (Any)
terms (Any)

Return type:

Any

class geojax.geometry.SPDBuresWasserstein(size, *, atol=1e-06, eps=1e-10)#

Bases: ExactGeometryMixin

Bures-Wasserstein geometry on SPD(size).

If P = Q diag(d) Q.T and U_tilde = Q.T @ U @ Q, the metric is

g_P(U,V) = 1/2 sum_ij U_tilde_ij V_tilde_ij / (d_i + d_j).

The exponential is defined only while its horizontal square-root lift is nonsingular. transport is an exact isometric vector transport for optimization; it is not the Levi-Civita parallel transport, whose general evaluation requires integrating a differential equation.

Parameters:

size (tuple[int, int])
atol (float)
eps (float)

sylvester(P, U)#

Solve P A + A P = U for symmetric A.

Parameters:

P (Any)
U (Any)

Return type:

Any

retr(P, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

P (Any)
U (Any)
t (float | Any)

Return type:

Any

optimal_transport_map(P, Q)#

Return the optimal Gaussian transport map from covariance P to Q.

Parameters:

P (Any)
Q (Any)

Return type:

Any

squared_dist(P, Q)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

P (Any)
Q (Any)

Return type:

Any

transport(P, Q, U)#

Isometric vector transport between Bures-Wasserstein tangent spaces.

This transport preserves the Bures-Wasserstein metric exactly but is not claimed to be Levi-Civita parallel transport.

Parameters:

P (Any)
Q (Any)
U (Any)

Return type:

Any

transp(P, Q, U)#

Isometric vector transport between Bures-Wasserstein tangent spaces.

This transport preserves the Bures-Wasserstein metric exactly but is not claimed to be Levi-Civita parallel transport.

Parameters:

P (Any)
Q (Any)
U (Any)

Return type:

Any

lincomb(P, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

P (Any)
terms (Any)

Return type:

Any

Fixed-rank and semidefinite matrices#

class geojax.geometry.FixedRank(size, *, rank, atol=1e-06, eps=1e-10)#

Bases: RetractionGeometryMixin

Embedded manifold of real matrices with fixed rank.

The Frobenius metric is used. retr is the truncated-SVD retraction; compatibility exp, log and dist calls are explicitly marked as proxies by operation_kind().

Parameters:

size (tuple[int, int])
rank (int)
atol (float)
eps (float)

retr(X, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

X (Any)
U (Any)
t (float | Any)

Return type:

Any

invretr(X, Y)#

Default inverse retraction: use the logarithmic map.

Parameters:

X (Any)
Y (Any)

Return type:

Any

dist(x, y)#

Local inverse-retraction norm; not a geodesic distance.

Parameters:

x (Any)
y (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp(x, u)#

Retraction proxy for the exponential map.

Parameters:

x (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log(x, y)#

Inverse-retraction proxy for the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

transp(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

transport(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

class geojax.geometry.RankKPSD(size, *, rank, atol=1e-06, eps=1e-10)#

Bases: _RankKPSDBase

Fixed-rank positive-semidefinite matrices with embedded metric.

Parameters:

size (tuple[int, int])
rank (int)
atol (float)
eps (float)

dist(x, y)#

Local inverse-retraction norm; not a geodesic distance.

Parameters:

x (Any)
y (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp(x, u)#

Retraction proxy for the exponential map.

Parameters:

x (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

invretr(P, Q)#

Default inverse retraction: use the logarithmic map.

Parameters:

P (Any)
Q (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log(x, y)#

Inverse-retraction proxy for the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

retr(P, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

P (Any)
U (Any)
t (float | Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

transp(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

transport(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

class geojax.geometry.RankKPSDBuresWasserstein(size, *, rank, atol=1e-06, eps=1e-10)#

Bases: _RankKPSDBase

Fixed-rank PSD matrices with their Bures–Wasserstein quotient metric.

Parameters:

size (tuple[int, int])
rank (int)
atol (float)
eps (float)

exp(P, U)#

Retraction proxy for the exponential map.

Parameters:

P (Any)
U (Any)

Return type:

Any

retr(P, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

P (Any)
U (Any)
t (float | Any)

Return type:

Any

log(P, Q)#

Inverse-retraction proxy for the logarithmic map.

Parameters:

P (Any)
Q (Any)

Return type:

Any

invretr(P, Q)#

Default inverse retraction: use the logarithmic map.

Parameters:

P (Any)
Q (Any)

Return type:

Any

dist(P, Q)#

Local inverse-retraction norm; not a geodesic distance.

Parameters:

P (Any)
Q (Any)

Return type:

Any

transport(P, Q, U)#

Projection vector transport associated with the retraction.

Parameters:

P (Any)
Q (Any)
U (Any)

Return type:

Any

transp(P, Q, U)#

Projection vector transport associated with the retraction.

Parameters:

P (Any)
Q (Any)
U (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

class geojax.geometry.Elliptope(size, *, rank, atol=1e-06, eps=1e-10)#

Bases: _RankKPSDBase

Rank-rank PSD matrices with unit diagonal.

Parameters:

size (tuple[int, int])
rank (int)
atol (float)
eps (float)

dist(x, y)#

Local inverse-retraction norm; not a geodesic distance.

Parameters:

x (Any)
y (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp(x, u)#

Retraction proxy for the exponential map.

Parameters:

x (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

invretr(P, Q)#

Default inverse retraction: use the logarithmic map.

Parameters:

P (Any)
Q (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log(x, y)#

Inverse-retraction proxy for the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

retr(P, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

P (Any)
U (Any)
t (float | Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

transp(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

transport(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

class geojax.geometry.Spectrahedron(size, *, rank, atol=1e-06, eps=1e-10)#

Bases: _RankKPSDBase

Rank-rank PSD matrices with unit trace.

Parameters:

size (tuple[int, int])
rank (int)
atol (float)
eps (float)

dist(x, y)#

Local inverse-retraction norm; not a geodesic distance.

Parameters:

x (Any)
y (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp(x, u)#

Retraction proxy for the exponential map.

Parameters:

x (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

invretr(P, Q)#

Default inverse retraction: use the logarithmic map.

Parameters:

P (Any)
Q (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log(x, y)#

Inverse-retraction proxy for the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

retr(P, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

P (Any)
U (Any)
t (float | Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

transp(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

transport(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

Correlation matrices#

Euclidean-Cholesky metric#

class geojax.geometry.CorrelationECM(size, *, atol=1e-06, eps=1e-10)#

Bases: _CorrelationCholeskyBase

Correlation manifold with the Euclidean-Cholesky metric.

ECM is the Euclidean-Cholesky metric. Its flat coordinates are the strictly lower-triangular entries of the unit-diagonal Cholesky factor.

Parameters:

size (tuple[int, int])
atol (float)
eps (float)

Log-Euclidean-Cholesky metric#

class geojax.geometry.CorrelationLEC(size, *, atol=1e-06, eps=1e-10)#

Bases: _CorrelationCholeskyBase

Correlation manifold with the log-Euclidean-Cholesky metric.

LEC is the log-Euclidean-Cholesky metric (also abbreviated LECM in the literature). Its flat coordinates are the matrix logarithm of the unit-diagonal Cholesky factor.

Parameters:

size (tuple[int, int])
atol (float)
eps (float)

Affine-invariant quotient metric#

class geojax.geometry.CorrelationAffineQuotient(size, *, atol=1e-06, eps=1e-10)#

Bases: RetractionGeometryMixin

Full-rank correlations with the affine-invariant quotient metric.

Correlation matrices are the quotient of SPD matrices by positive diagonal congruence. The metric is evaluated using affine-invariant horizontal lifts. Point updates use normalized-addition retractions; consequently exp, log and dist are documented proxies.

Parameters:

size (tuple[int, int])
atol (float)
eps (float)

horizontal_lift(C, U)#

Lift a correlation tangent horizontally to the SPD total space.

Parameters:

C (Any)
U (Any)

Return type:

Any

retr(C, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

C (Any)
U (Any)
t (float | Any)

Return type:

Any

invretr(C, D)#

Default inverse retraction: use the logarithmic map.

Parameters:

C (Any)
D (Any)

Return type:

Any

dist(x, y)#

Local inverse-retraction norm; not a geodesic distance.

Parameters:

x (Any)
y (Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp(x, u)#

Retraction proxy for the exponential map.

Parameters:

x (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log(x, y)#

Inverse-retraction proxy for the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

transp(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

transport(x, y, u)#

Projection vector transport associated with the retraction.

Parameters:

x (Any)
y (Any)
u (Any)

Return type:

Any

Shape spaces#

class geojax.geometry.KendallShape(size, *, atol=1e-06, eps=1e-10)#

Bases: ExactGeometryMixin

Regular Kendall shape space of centered, scale-normalized landmarks.

Public points are pre-shape matrices of size landmarks x ambient_dim. Frames related by a right action of SO(ambient_dim) represent the same shape. Tangents are represented by horizontal pre-shape vectors.

Parameters:

size (tuple[int, int])
atol (float)
eps (float)

align(Y, X)#

Align Y to X by orientation-preserving Procrustes rotation.

Parameters:

Y (Any)
X (Any)

Return type:

tuple[Any, Any]

retr(X, U, t=1.0)#

Default retraction: use the exponential map.

Parameters:

X (Any)
U (Any)
t (float | Any)

Return type:

Any

dist_batch(x, ys)#

Compatibility wrapper for natively batched dist.

Parameters:

x (Any)
ys (Any)

Return type:

Any

ehess_to_rhess(x, egrad, ehess_vec, u)#

Default Hessian conversion: tangent-project the ambient Hessian-vector product.

Parameters:

x (Any)
egrad (Any)
ehess_vec (Any)
u (Any)

Return type:

Any

exp_batch(x, us)#

Compatibility wrapper for natively batched exp.

Parameters:

x (Any)
us (Any)

Return type:

Any

invretr(x, y)#

Default inverse retraction: use the logarithmic map.

Parameters:

x (Any)
y (Any)

Return type:

Any

lincomb(x, *terms)#

Linear combination of tangent vectors, projected back to T_x M.

Parameters:

x (Any)
terms (Any)

Return type:

Any

log_batch(x, ys)#

Compatibility wrapper for natively batched log.

Parameters:

x (Any)
ys (Any)

Return type:

Any

operation_kind(name)#

Describe the mathematical status of a geometric operation.

Parameters:: name (str)
Return type:: str

pair_mean(x, y)#

Midpoint-like construction from the available exp and log.

This is a geodesic midpoint only when both operations are exact and the selected logarithm is the unique minimizing one.

Parameters:

x (Any)
y (Any)

Return type:

Any

squared_dist(x, y)#

Squared geodesic distance, evaluated without differentiating sqrt.

Parameters:

x (Any)
y (Any)

Return type:

Any

Utility functions#

geojax.geometry.torus.wrap_angles(x)#

Wrap angles to [-pi, pi).

Parameters:: x (Any)
Return type:: Any