GeoJAX

GeoJAX is a JAX-native toolkit for Riemannian geometry, manifold optimization, and statistical and machine learning with manifold-valued data. Geometries define representations and exact primitives where available; optimization and learning methods consume that common protocol without flattening every scientific object into an ordinary vector.

import jax

from geojax.geometry import Sphere
from geojax.learning import as_manifold_data, frechet_mean, pairwise_distances

M = Sphere(size=3)
observations = M.random_point(jax.random.key(0), sample_shape=(32,))
data = as_manifold_data(M, observations)

center = frechet_mean(M, data).point
distances = pairwise_distances(M, data)

Getting started

Install GeoJAX, validate manifold data, and solve a first geometric learning problem.

Tutorials

Work through executable mathematical examples with computed figures.

User guide

Understand geometry conventions, learning workflows, and solver choices.

API reference

Look up public geometries, manifold-learning methods, optimization problems, and solvers.

Scope

GeoJAX is alpha software. Geometry capabilities, derivative contracts, and learning-data validation are explicit. Optimization and learning methods are tested on analytic cases, representative manifolds, and executable tutorials.