{"data":{"slug":"jax-ml-jax","name":"jax","tagline":"Composable transformations of Python+NumPy programs","github_url":"https://github.com/jax-ml/jax","owner":"jax-ml","repo":"jax","owner_avatar_url":"https://avatars.githubusercontent.com/u/58486408?v=4","primary_language":"Python","stars":36338,"forks":3810,"topics":["jax"],"archived":false,"github_pushed_at":"2026-09-26T05:38:44+00:00","maintenance_label":"Very active","stars_delta_30d":253,"url":"https://www.graphcanon.com/tools/jax-ml-jax","markdown_url":"https://www.graphcanon.com/tools/jax-ml-jax.md","api_url":"https://www.graphcanon.com/api/graphcanon/tools/jax-ml-jax","graph_url":"https://www.graphcanon.com/api/graphcanon/graph?tool=jax-ml-jax","description":"Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more","homepage_url":"https://docs.jax.dev","license":"Apache-2.0","open_issues":2597,"watchers":333,"ai_summary":"JAX is a Python library for high-performance numerical computing and machine learning on accelerators such as GPUs and TPUs, supporting automatic differentiation and compilation.","readme_excerpt":"<div align=\"center\">\n<img src=\"https://raw.githubusercontent.com/jax-ml/jax/main/images/jax_logo_250px.png\" alt=\"logo\"></img>\n</div>\n\n# Transformable numerical computing at scale\n\n\n\n\n[**Transformations**](#transformations)\n| [**Scaling**](#scaling)\n| [**Install guide**](#installation)\n| [**Change logs**](https://docs.jax.dev/en/latest/changelog.html)\n| [**Reference docs**](https://docs.jax.dev/en/latest/)\n\n\n## What is JAX?\n\nJAX is a Python library for accelerator-oriented array computation and program transformation,\ndesigned for high-performance numerical computing and large-scale machine learning.\n\nJAX can automatically differentiate native\nPython and NumPy functions. It can differentiate through loops, branches,\nrecursion, and closures, and it can take derivatives of derivatives of\nderivatives. It supports reverse-mode differentiation (a.k.a. backpropagation)\nvia [`jax.grad`](#automatic-differentiation-with-grad) as well as forward-mode differentiation,\nand the two can be composed arbitrarily to any order.\n\nJAX uses [XLA](https://www.openxla.org/xla)\nto compile and scale your NumPy programs on TPUs, GPUs, and other hardware accelerators.\nYou can compile your own pure functions with [`jax.jit`](#compilation-with-jit).\nCompilation and automatic differentiation can be composed arbitrarily.\n\nDig a little deeper, and you'll see that JAX is really an extensible system for\n[composable function transformations](#transformations) at [scale](#scaling).\n\nThis is a research project, not an official Google product. Expect\n[sharp edges](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html).\nPlease help by trying it out, [reporting bugs](https://github.com/jax-ml/jax/issues),\nand letting us know what you think!\n\n```python\nimport jax\nimport jax.numpy as jnp\n\ndef predict(params, inputs):\n  for W, b in params:\n    outputs = jnp.dot(inputs, W) + b\n    inputs = jnp.tanh(outputs)  # inputs to the next layer\n  return outputs                # no activation on last layer\n\ndef loss(params, inputs, targets):\n  preds = predict(params, inputs)\n  return jnp.sum((preds - targets)**2)\n\ngrad_loss = jax.jit(jax.grad(loss))  # compiled gradient evaluation function\nperex_grads = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))  # fast per-example grads\n```\n\n### Contents\n* [Transformations](#transformations)\n* [Scaling](#scaling)\n* [Current gotchas](#gotchas-and-sharp-bits)\n* [Installation](#installation)\n* [Citing JAX](#citing-jax)\n* [Reference documentation](#reference-documentation)\n\n## Transformations\n\nAt its core, JAX is an extensible system for transforming numerical functions.\nHere are three: `jax.grad`, `jax.jit`, and `jax.vmap`.\n\n### Automatic differentiation with `grad`\n\nUse [`jax.grad`](https://docs.jax.dev/en/latest/jax.html#jax.grad)\nto efficiently compute reverse-mode gradients:\n\n```python\nimport jax\nimport jax.numpy as jnp\n\ndef tanh(x):\n  y = jnp.exp(-2.0 * x)\n  return (1.0 - y) / (1.0 + y)\n\ngrad_tanh = jax.grad(tanh)\nprint(grad_tanh(1.0))\n# prints 0.4199743\n```\n\nYou can differentiate to any order with `grad`:\n\n```python\nprint(jax.grad(jax.grad(jax.grad(tanh)))(1.0))\n# prints 0.62162673\n```\n\nYou're free to use differentiation with Python control flow:\n\n```python\ndef abs_val(x):\n  if x > 0:\n    return x\n  else:\n    return -x\n\nabs_val_grad = jax.grad(abs_val)\nprint(abs_val_grad(1.0))   # prints 1.0\nprint(abs_val_grad(-1.0))  # prints -1.0 (abs_val is re-evaluated)\n```\n\nSee the [JAX Autodiff\nCookbook](https://docs.jax.dev/en/latest/notebooks/autodiff_cookbook.html)\nand the [reference docs on automatic\ndifferentiation](https://docs.jax.dev/en/latest/jax.html#automatic-differentiation)\nfor more.\n\n### Compilation with `jit`\n\nUse XLA to compile your functions end-to-end with\n[`jit`](https://docs.jax.dev/en/latest/jax.html#just-in-time-compilation-jit),\nused either as an `@jit` decorator or as a higher-order function.\n\n```python\nimport jax\nimport jax.numpy as jnp\n\ndef slow_f(x):\n  # Element-wise ops see a large benefit from fusion\n  re","github_created_at":"2018-10-25T21:25:02+00:00","created_at":"2026-07-11T23:23:16.062176+00:00","updated_at":"2026-09-26T06:00:35.211472+00:00","categories":[{"slug":"inference-serving","name":"Inference & Serving","url":"https://www.graphcanon.com/categories/inference-serving","markdown_url":"https://www.graphcanon.com/categories/inference-serving.md","api_url":"https://www.graphcanon.com/api/graphcanon/categories/inference-serving"},{"slug":"model-training","name":"Model Training","url":"https://www.graphcanon.com/categories/model-training","markdown_url":"https://www.graphcanon.com/categories/model-training.md","api_url":"https://www.graphcanon.com/api/graphcanon/categories/model-training"}],"tags":[{"slug":"compilation","name":"compilation"},{"slug":"differentiation","name":"differentiation"},{"slug":"gpu","name":"gpu"},{"slug":"python","name":"python"},{"slug":"tpu","name":"tpu"}],"trust":{"provenance":{"is_fork":false,"github_id":154739597,"owner_type":"Organization","methodology":"github_public_v1","parent_repo":null,"near_duplicate_slugs":[]},"computed_at":"2026-09-26T06:00:34.427Z","maintenance":{"label":"Very active","score":96,"methodology":"github_public_v1","releases_90d":3,"days_since_push":0,"last_release_at":"2026-09-17T23:43:51Z","stars_delta_30d":253,"open_issues_delta_30d":52},"security_summary":{"status":"no_lockfile","scanner":null,"low_count":0,"high_count":0,"last_scan_at":"2026-07-11T23:23:20.544Z","medium_count":0,"scan_profile":"none","critical_count":0}},"capability_facts":{"scan":{"source":"repo_scan","observed_at":"2026-09-26T06:00:34.881Z"},"languages":{"value":["python"],"source":"github.language+pyproject.toml","observed_at":"2026-09-26T06:00:34.881Z"},"license_spdx":{"value":"Apache-2.0","source":"github.license","observed_at":"2026-09-26T06:00:34.881Z"}},"decision_facts":{"hosting":null,"pricing":null,"requirements":null,"constraints":null,"when_to_use":["- When you need to perform high-performance numerical computations with support for both forward and reverse mode automatic differentiation on accelerators such as GPUs and TPUs.","- If your project requires efficient handling of control flow constructs like loops and branches during the differentiation process, thus ensuring smooth integration with complex Python functions."],"when_not_to_use":["- JAX should be avoided if your codebase heavily relies on non-JIT compatible operations or side effects within Python functions, due to JAX's limitations in those areas.","- For applications that do not require GPU/TPU acceleration and where performance gains from automatic differentiation and compilation are not critical."],"source":"enrich:decision_facts","observed_at":"2026-07-14T18:49:03.195Z"},"constraint_facets":null,"decision_summary":[{"label":"Adopt for","value":"JAX is a high-performance numerical computing library for Python that integrates automatic differentiation and compilation, suitable for GPU and TPU acceleration."}]}}