{"data":{"slug":"jax-ml-jax","name":"jax","tagline":"Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more","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":35999,"forks":3676,"topics":["jax"],"archived":false,"github_pushed_at":"2026-07-11T17:29:35+00:00","maintenance_label":"Very active","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":2495,"watchers":329,"ai_summary":null,"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-07-11T23:23:26.36206+00:00","categories":[{"slug":"vector-databases","name":"Vector Databases","url":"https://www.graphcanon.com/categories/vector-databases","markdown_url":"https://www.graphcanon.com/categories/vector-databases.md","api_url":"https://www.graphcanon.com/api/graphcanon/categories/vector-databases"},{"slug":"computer-vision","name":"Computer Vision","url":"https://www.graphcanon.com/categories/computer-vision","markdown_url":"https://www.graphcanon.com/categories/computer-vision.md","api_url":"https://www.graphcanon.com/api/graphcanon/categories/computer-vision"},{"slug":"evaluation-observability","name":"Evaluation & Observability","url":"https://www.graphcanon.com/categories/evaluation-observability","markdown_url":"https://www.graphcanon.com/categories/evaluation-observability.md","api_url":"https://www.graphcanon.com/api/graphcanon/categories/evaluation-observability"}],"tags":[{"slug":"python","name":"python"},{"slug":"jax","name":"jax"}],"trust":{"provenance":{"is_fork":false,"github_id":154739597,"owner_type":"Organization","methodology":"github_public_v1","parent_repo":null,"near_duplicate_slugs":[]},"computed_at":"2026-07-11T23:23:20.029Z","maintenance":{"label":"Very active","score":96,"methodology":"github_public_v1","releases_90d":3,"days_since_push":0,"last_release_at":"2026-06-17T23:39:56Z"},"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-07-11T23:23:19.789Z"},"languages":{"value":["python"],"source":"github.language+pyproject.toml","observed_at":"2026-07-11T23:23:19.789Z"},"license_spdx":{"value":"Apache-2.0","source":"github.license","observed_at":"2026-07-11T23:23:19.789Z"}}}}