{"id":17191062,"url":"https://github.com/dfm/jpu","last_synced_at":"2025-04-13T19:31:29.096Z","repository":{"id":39072627,"uuid":"495986012","full_name":"dfm/jpu","owner":"dfm","description":"JAX + Units","archived":false,"fork":false,"pushed_at":"2025-02-01T18:24:44.000Z","size":99,"stargazers_count":47,"open_issues_count":3,"forks_count":2,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-04-12T06:38:30.124Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","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/dfm.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","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":"2022-05-24T20:57:57.000Z","updated_at":"2025-03-08T01:09:05.000Z","dependencies_parsed_at":"2023-12-20T15:26:52.770Z","dependency_job_id":"d99d067f-205d-49ac-87cf-1fc0c357afbe","html_url":"https://github.com/dfm/jpu","commit_stats":null,"previous_names":[],"tags_count":6,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dfm%2Fjpu","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dfm%2Fjpu/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dfm%2Fjpu/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dfm%2Fjpu/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/dfm","download_url":"https://codeload.github.com/dfm/jpu/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248767897,"owners_count":21158555,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","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":[],"created_at":"2024-10-15T01:24:46.953Z","updated_at":"2025-04-13T19:31:28.793Z","avatar_url":"https://github.com/dfm.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# JAX + Units\r\n\r\n**Built with [JAX](https://jax.readthedocs.io) and\r\n[Pint](https://pint.readthedocs.io)!**\r\n\r\nThis module provides an interface between [JAX](https://jax.readthedocs.io) and\r\n[Pint](https://pint.readthedocs.io) to allow JAX to support operations with\r\nunits. The propagation of units happens at trace time, so jitted functions\r\nshould see no runtime cost. This library is experimental so expect some sharp\r\nedges.\r\n\r\nFor example:\r\n\r\n```python\r\n\u003e\u003e\u003e import jax\r\n\u003e\u003e\u003e import jax.numpy as jnp\r\n\u003e\u003e\u003e import jpu\r\n\u003e\u003e\u003e\r\n\u003e\u003e\u003e u = jpu.UnitRegistry()\r\n\u003e\u003e\u003e\r\n\u003e\u003e\u003e @jax.jit\r\n... def add_two_lengths(a, b):\r\n...     return a + b\r\n...\r\n\u003e\u003e\u003e add_two_lengths(3 * u.m, jnp.array([4.5, 1.2, 3.9]) * u.cm)\r\n\u003cQuantity([3.045 3.012 3.039], 'meter')\u003e\r\n\r\n```\r\n\r\n## Installation\r\n\r\nTo install, use `pip`:\r\n\r\n```bash\r\npython -m pip install jpu\r\n```\r\n\r\nThe only dependencies are `jax` and `pint`, and these will also be installed, if\r\nnot already in your environment. Take a look at [the JAX docs for more\r\ninformation about installing JAX on different\r\nsystems](https://github.com/google/jax#installation).\r\n\r\n## Usage\r\n\r\nHere is a slightly more complete example:\r\n\r\n```python\r\n\u003e\u003e\u003e import jax\r\n\u003e\u003e\u003e import numpy as np\r\n\u003e\u003e\u003e from jpu import UnitRegistry, numpy as jnpu\r\n\u003e\u003e\u003e\r\n\u003e\u003e\u003e u = UnitRegistry()\r\n\u003e\u003e\u003e\r\n\u003e\u003e\u003e @jax.jit\r\n... def projectile_motion(v_init, theta, time, g=u.standard_gravity):\r\n...     \"\"\"Compute the motion of a projectile with support for units\"\"\"\r\n...     x = v_init * time * jnpu.cos(theta)\r\n...     y = v_init * time * jnpu.sin(theta) - 0.5 * g * jnpu.square(time)\r\n...     return x.to(u.m), y.to(u.m)\r\n...\r\n\u003e\u003e\u003e x, y = projectile_motion(\r\n...     5.0 * u.km / u.h, 60 * u.deg, np.linspace(0, 1, 50) * u.s\r\n... )\r\n\r\n```\r\n\r\n## Technical details \u0026 limitations\r\n\r\nThe most significant limitation of this library is the fact that users must use\r\n`jpu.numpy` functions when interacting with \"quantities\" with units instead of\r\nthe `jax.numpy` interface. This is because JAX does not (yet?) provide a general\r\ninterface for dispatching of ufuncs on custom array classes. I have played\r\naround with the undocumented `__jax_array__` interface, but it's not really\r\nflexible enough, and it isn't currently compatible with Pytree objects.\r\n\r\nSo far, only a subset of the `numpy`/`jax.numpy` interface is implemented. Pull\r\nrequests adding broader support (including submodules) would be welcome!\r\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdfm%2Fjpu","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdfm%2Fjpu","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdfm%2Fjpu/lists"}