{"id":37724703,"url":"https://github.com/danielkelshaw/riemax","last_synced_at":"2026-01-18T00:29:31.228Z","repository":{"id":203973077,"uuid":"706387319","full_name":"danielkelshaw/riemax","owner":"danielkelshaw","description":"Riemannian geometry in JAX","archived":false,"fork":false,"pushed_at":"2024-08-21T00:53:16.000Z","size":2432,"stargazers_count":13,"open_issues_count":0,"forks_count":2,"subscribers_count":1,"default_branch":"main","last_synced_at":"2024-12-06T21:45:11.053Z","etag":null,"topics":["deep-learning","differential-geometry","geodesics","jax","riemannian-geometry"],"latest_commit_sha":null,"homepage":"https://riemax.danielkelshaw.com","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/danielkelshaw.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-10-17T21:07:35.000Z","updated_at":"2024-10-09T22:23:45.000Z","dependencies_parsed_at":"2023-11-14T16:43:21.051Z","dependency_job_id":"2d70958c-04ee-4bf7-a516-df6d8c34ec90","html_url":"https://github.com/danielkelshaw/riemax","commit_stats":null,"previous_names":["danielkelshaw/riemax"],"tags_count":2,"template":false,"template_full_name":null,"purl":"pkg:github/danielkelshaw/riemax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/danielkelshaw%2Friemax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/danielkelshaw%2Friemax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/danielkelshaw%2Friemax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/danielkelshaw%2Friemax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/danielkelshaw","download_url":"https://codeload.github.com/danielkelshaw/riemax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/danielkelshaw%2Friemax/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":28479033,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-01-16T11:59:17.896Z","status":"ssl_error","status_checked_at":"2026-01-16T11:55:55.838Z","response_time":107,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","differential-geometry","geodesics","jax","riemannian-geometry"],"created_at":"2026-01-16T13:38:24.453Z","updated_at":"2026-01-16T13:38:25.192Z","avatar_url":"https://github.com/danielkelshaw.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Riemax: Riemannian geometry for JAX\n\n\u003e [!IMPORTANT]\n\u003e This library is very much a work in progress, as is the documentation. Naturally, the contents of this library is subject to change depending on choice of research avenue. At the moment, there exists a solid basis on which to build, but the API may change to improve accessiblity for new users. If there is anything you would like to see in the library, or anything you think that needs to change -- please do let me know.\n\nRiemax is a [JAX] library for Riemannian geometry, providing the ability to define induced metrics and operate on manifolds directly. This includes functionality, such as:\n\n  - Computing geometric quantities on manifolds.\n  - Defining operators on manifolds.\n  - Tools for geometric statistics.\n\n## Installation\n\n```bash\npip install git+https://github.com/danielkelshaw/riemax\n```\n\nRequires Python 3.12+ and JAX 0.4.13+.\n\n## Quick Example\n\nManifolds can be defined by a function transformation, $\\iota: M \\rightarrow N$:\n\n```python\nimport riemax as rx\n\nfn_transformation = rx.fn_transformations.fn_peaks\nmanifold = rx.Manifold.from_fn_transformation(fn_transformation)\n```\n\nand can be used to compute properties on the manifold:\n\n```python\nimport jax\nimport jax.numpy as jnp\n\np = jnp.array([0.0, 0.0])\n\nmetric = manifold.metric_tensor(p)\nchristoffel = manifold.sk_christoffel(p)\nricci_scalar = manifold.ricci_scalar(p)\n```\nWe can also define operators. Given a function $f: M \\rightarrow \\mathbb{R}$:\n\n```python\nfrom riemax.manifold import M, Rn\n\ndef f(p: M[jax.Array]) -\u003e Rn[jax.Array]:\n    return jnp.einsum('i -\u003e ', p ** 2)\n\nfn_grad = manifold.grad(f)\n\n# we can define the laplacian explicitly:\nfn_laplacian = manifold.div(fn_grad)\n\n# ... or from manifold:\nfn_laplacian = manifold.laplace_beltrami(f)\n```\n\nWe can exploit the ability to compute geometric quantities to compute the exponential map:\n\n```python\ndt = 1e-3\nn_steps = int(1.0 // dt)\n\nfn_exp_map = rx.manifold.maps.exponential_map_factory(\n    integrator=rx.numerical.integrators.odeint,\n    dt=dt,\n    metric=manifold.metric_tensor,\n    n_steps=n_steps\n)\n\np_in_tm = rx.manifold.TangentSpace(p, jnp.array([1.0, 1.0]))\nq_in_tm, trajectory = fn_exp_map(p_in_tm)\n```\n\n## Next Steps\nThe examples above show just a fraction of what is possible with Riemax. If this quick start has piqued your interest, please feel free to take a look at the examples to get a better feeling for what is possible with this library.\n\n## Citation\n\nIf you found this library to be useful in academic work, then please cite:\n\n```tex\n@misc{kelshaw2023riemax\n    title = {{Riemax}: {R}iemannian geometry in {JAX} via automatic differentiation}\n    author = {Daniel Kelshaw}\n    year = {2023},\n    url = {https://github.com/danielkelshaw/riemax}\n}\n```\n\n\n[JAX]: https://github.com/google/jax\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdanielkelshaw%2Friemax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdanielkelshaw%2Friemax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdanielkelshaw%2Friemax/lists"}