{"id":13472546,"url":"https://github.com/patrick-kidger/jaxtyping","last_synced_at":"2026-01-05T12:21:45.814Z","repository":{"id":44475773,"uuid":"506654025","full_name":"patrick-kidger/jaxtyping","owner":"patrick-kidger","description":"Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/","archived":false,"fork":false,"pushed_at":"2024-10-25T19:11:15.000Z","size":299,"stargazers_count":1168,"open_issues_count":42,"forks_count":62,"subscribers_count":13,"default_branch":"main","last_synced_at":"2024-10-29T15:38:57.310Z","etag":null,"topics":["jax","python-typing","typing"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/patrick-kidger.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","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":"2022-06-23T13:35:25.000Z","updated_at":"2024-10-29T09:40:27.000Z","dependencies_parsed_at":"2023-02-12T18:00:29.696Z","dependency_job_id":"a9c09482-d446-41c5-bb25-d3d0bd4f6e10","html_url":"https://github.com/patrick-kidger/jaxtyping","commit_stats":{"total_commits":154,"total_committers":14,"mean_commits":11.0,"dds":"0.11688311688311692","last_synced_commit":"07e58de0bdbc75087a77575a6718963802c84a36"},"previous_names":["google/jaxtyping"],"tags_count":38,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Fjaxtyping","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Fjaxtyping/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Fjaxtyping/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/patrick-kidger%2Fjaxtyping/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/patrick-kidger","download_url":"https://codeload.github.com/patrick-kidger/jaxtyping/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248688543,"owners_count":21145764,"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":["jax","python-typing","typing"],"created_at":"2024-07-31T16:00:55.622Z","updated_at":"2026-01-05T12:21:45.807Z","avatar_url":"https://github.com/patrick-kidger.png","language":"Python","readme":"\u003ch1 align=\"center\"\u003ejaxtyping\u003c/h1\u003e\n\nA library providing type annotations **and runtime type-checking** for the shape and dtype of JAX/PyTorch/NumPy/MLX/TensorFlow arrays and tensors.\n\n_The name 'jax'typing is now historical, we support all of the above and have no JAX dependency!_\n\n```python\nfrom jaxtyping import Float\nfrom torch import Tensor\n\n# Accepts floating-point 2D arrays with matching axes\ndef matrix_multiply(x: Float[Tensor, \"dim1 dim2\"],\n                    y: Float[Tensor, \"dim2 dim3\"]\n                  ) -\u003e Float[Tensor, \"dim1 dim3\"]:\n    ...\n```\n\n## Installation\n\n```bash\npip install jaxtyping\n```\n\nRequires Python 3.10+.\n\nThe annotations provided by jaxtyping are compatible with runtime type-checking packages, so it is common to also install one of these. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments).\n\n## Documentation\n\nAvailable at [https://docs.kidger.site/jaxtyping](https://docs.kidger.site/jaxtyping).\n\n## See also: other libraries in the JAX ecosystem\n\n**Always useful**  \n[Equinox](https://github.com/patrick-kidger/equinox): neural networks and everything not already in core JAX!  \n\n**Deep learning**  \n[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.  \n[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).  \n[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).  \n[paramax](https://github.com/danielward27/paramax): parameterizations and constraints for PyTrees.  \n\n**Scientific computing**  \n[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.  \n[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.  \n[Lineax](https://github.com/patrick-kidger/lineax): linear solvers.  \n[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.  \n[sympy2jax](https://github.com/patrick-kidger/sympy2jax): SymPy\u003c-\u003eJAX conversion; train symbolic expressions via gradient descent.  \n[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)  \n\n**Awesome JAX**  \n[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of other JAX projects.  \n","funding_links":[],"categories":["Python","See also: other libraries in the JAX ecosystem","Libraries"],"sub_categories":[],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpatrick-kidger%2Fjaxtyping","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpatrick-kidger%2Fjaxtyping","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpatrick-kidger%2Fjaxtyping/lists"}