False
        

    Skip to content

      Navigation Menu

            Sign in Appearance settings
                * Platform
                        + AI CODE CREATION
                            o GitHub Copilot Write better code with AI
                            o GitHub Copilot app Direct agents from issue to merge
                            o MCP Registry Integrate external tools
                        + DEVELOPER WORKFLOWS
                            o Actions Automate any workflow
                            o Codespaces Instant dev environments
                            o Issues Plan and track work
                            o Code Review Manage code changes
                            o Code Quality Enforce quality at merge
                        + APPLICATION SECURITY
                            o GitHub Advanced Security Find and fix vulnerabilities
                            o Code security Secure your code as you build
                            o Secret protection Stop leaks before they start
                        + EXPLORE
                            o Why GitHub
                            o Documentation
                            o Blog
                            o Changelog
                            o Marketplace
                      View all features
                * Solutions
                        + BY COMPANY SIZE
                            o Enterprises
                            o Small and medium teams
                            o Startups
                            o Nonprofits
                        + BY USE CASE
                            o App Modernization
                            o DevSecOps
                            o DevOps
                            o CI/CD
                            o View all use cases
                        + BY INDUSTRY
                            o Healthcare
                            o Financial services
                            o Manufacturing
                            o Government
                            o View all industries
                      View all solutions
                * Resources
                        + EXPLORE BY TOPIC
                            o AI
                            o Software Development
                            o DevOps
                            o Security
                            o View all topics
                        + EXPLORE BY TYPE
                            o Customer stories
                            o Events & webinars
                            o Ebooks & reports
                            o Business insights
                            o GitHub Skills
                        + SUPPORT & SERVICES
                            o Documentation
                            o Customer support
                            o Community forum
                            o Trust center
                            o Partners
                      View all resources
                * Open Source
                        + COMMUNITY
                            o GitHub Sponsors Fund open source developers
                        + PROGRAMS
                            o Security Lab
                            o Maintainer Community
                            o Accelerator
                            o GitHub Stars
                            o Archive Program
                        + REPOSITORIES
                            o Topics
                            o Trending
                            o Collections
                * Enterprise
                        + ENTERPRISE SOLUTIONS
                            o Enterprise platform AI-powered developer platform
                        + AVAILABLE ADD-ONS
                            o GitHub Advanced Security Enterprise-grade security features
                            o Copilot for Business Enterprise-grade AI features
                            o Premium Support Enterprise-grade 24/7 support
              * Pricing
                Type / to search
                Sign in
              Sign up Appearance settings
      You signed in with another tab or window. Reload to refresh your session. You signed out in another tab or window. Reload to refresh your session. You switched accounts on another tab or window. Reload to refresh your session. Dismiss alert

              Uh oh!

              There was an error while loading. Please reload this page.

              jax-ml / jax Public
              * Notifications You must be signed in to change notification settings
              * Fork 3.7k
                * Star 36.1k
          * Code
          * Issues 1.7k
          * Pull requests 859
          * Pull requests 860
          * Discussions
          * Actions
          * Security and quality 0
          * Insights
          Additional navigation options
                  * Code
                  * Issues
                  * Pull requests
                  * Discussions
                  * Actions
                  * Security and quality
                  * Insights
                                              main
                                          Branches Tags
                                            Go to file
                                        Code
                                          Open more actions menu

                                        Folders and files

                                          Name                             Name                               Last commit message    Last commit date
                                            Latest commit

                                                History

                                                37,557 Commits
                                                  37,557 Commits
                                                37,603 Commits
                                                  37,603 Commits
                                                  .github                          .github
                                                  benchmarks                       benchmarks
                                                  build                            build
                                                  ci                               ci
                                                  cloud_tpu_colabs                 cloud_tpu_colabs
                                                  docs                             docs
                                                  examples                         examples
                                                  images                           images
                                                  jax                              jax
                                                  jax_plugins                      jax_plugins
                                                  jaxlib                           jaxlib
                                                  tests                            tests
                                                  third_party                      third_party
                                                  .bazelrc                         .bazelrc
                                                  .bazelversion                    .bazelversion
                                                  .editorconfig                    .editorconfig
                                                  .gitignore                       .gitignore
                                                  .pre-commit-config.yaml          .pre-commit-config.yaml
                                                  .readthedocs.yml                 .readthedocs.yml
                                                  AUTHORS                          AUTHORS
                                                  BUILD.bazel                      BUILD.bazel
                                                  CHANGELOG.md                     CHANGELOG.md
                                                  CITATION.bib                     CITATION.bib
                                                  CONTRIBUTING.md                  CONTRIBUTING.md
                                                  LICENSE                          LICENSE
                                                  MODULE.bazel                     MODULE.bazel
                                                  README.md                        README.md
                                                  WORKSPACE                        WORKSPACE
                                                  bazel_downloader.cfg             bazel_downloader.cfg
                                                  build_wheel.py                   build_wheel.py
                                                  conftest.py                      conftest.py
                                                  platform_mappings                platform_mappings
                                                  pyproject.toml                   pyproject.toml
                                                  setup.py                         setup.py
                                                  test_shard_count.bzl             test_shard_count.bzl
                                            View all files


                                              Repository files navigation

                                                *
                                                * README
                                                * Contributing
                                                * Apache-2.0 license
                                                More items

                                                Transformable numerical computing at scale

                                              Transformations | Scaling | Install guide | Change logs | Reference docs

                                                What is JAX?

                                              JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.

                                              JAX can automatically differentiate native Python and NumPy functions. It can differentiate through loops, branches, recursion, and closures, and it can take derivatives of derivatives of derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation) via jax.grad as well as forward-mode differentiation, and the two can be composed arbitrarily to any order.

                                              JAX uses XLA to compile and scale your NumPy programs on TPUs, GPUs, and other hardware accelerators. You can compile your own pure functions with jax.jit. Compilation and automatic differentiation can be composed arbitrarily.

                                              Dig a little deeper, and you'll see that JAX is really an extensible system for composable function transformations at scale.

                                              This is a research project, not an official Google product. Expect sharp edges. Please help by trying it out, reporting bugs, and letting us know what you think!

                                                import jax
                                                import jax.numpy as jnp

                                                def predict(params, inputs):
                                                  for W, b in params:
                                                    outputs = jnp.dot(inputs, W) + b
                                                    inputs = jnp.tanh(outputs)  # inputs to the next layer
                                                  return outputs                # no activation on last layer

                                                def loss(params, inputs, targets):
                                                  preds = predict(params, inputs)
                                                  return jnp.sum((preds - targets)**2)

                                                grad_loss = jax.jit(jax.grad(loss))  # compiled gradient evaluation function
                                                perex_grads = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))  # fast per-example grads

                                                Contents

                                                * Transformations
                                                * Scaling
                                                * Current gotchas
                                                * Installation
                                                * Citing JAX
                                                * Reference documentation

                                                Transformations

                                              At its core, JAX is an extensible system for transforming numerical functions. Here are three: jax.grad, jax.jit, and jax.vmap.

                                                Automatic differentiation with grad

                                              Use jax.grad to efficiently compute reverse-mode gradients:

                                                import jax
                                                import jax.numpy as jnp

                                                def tanh(x):
                                                  y = jnp.exp(-2.0 * x)
                                                  return (1.0 - y) / (1.0 + y)

                                                grad_tanh = jax.grad(tanh)
                                                print(grad_tanh(1.0))
                                                # prints 0.4199743

                                              You can differentiate to any order with grad:

                                                print(jax.grad(jax.grad(jax.grad(tanh)))(1.0))
                                                # prints 0.62162673

                                              You're free to use differentiation with Python control flow:

                                                def abs_val(x):
                                                  if x > 0:
                                                    return x
                                                  else:
                                                    return -x

                                                abs_val_grad = jax.grad(abs_val)
                                                print(abs_val_grad(1.0))   # prints 1.0
                                                print(abs_val_grad(-1.0))  # prints -1.0 (abs_val is re-evaluated)

                                              See the JAX Autodiff Cookbook and the reference docs on automatic differentiation for more.

                                                Compilation with jit

                                              Use XLA to compile your functions end-to-end with jit, used either as an @jit decorator or as a higher-order function.

                                                import jax
                                                import jax.numpy as jnp

                                                def slow_f(x):
                                                  # Element-wise ops see a large benefit from fusion
                                                  return x * x + x * 2.0

                                                x = jnp.ones((5000, 5000))
                                                fast_f = jax.jit(slow_f)
                                                %timeit -n10 -r3 fast_f(x)
                                                %timeit -n10 -r3 slow_f(x)

                                              Using jax.jit constrains the kind of Python control flow the function can use; see the tutorial on Control Flow and Logical Operators with JIT for more.

                                                Auto-vectorization with vmap

                                              vmap maps a function along array axes. But instead of just looping over function applications, it pushes the loop down onto the function’s primitive operations, e.g. turning matrix-vector multiplies into matrix-matrix multiplies for better performance.

                                              Using vmap can save you from having to carry around batch dimensions in your code:

                                                import jax
                                                import jax.numpy as jnp

                                                def l1_distance(x, y):
                                                  assert x.ndim == y.ndim == 1  # only works on 1D inputs
                                                  return jnp.sum(jnp.abs(x - y))

                                                def pairwise_distances(dist1D, xs):
                                                  return jax.vmap(jax.vmap(dist1D, (0, None)), (None, 0))(xs, xs)

                                                xs = jax.random.normal(jax.random.key(0), (100, 3))
                                                dists = pairwise_distances(l1_distance, xs)
                                                dists.shape  # (100, 100)

                                              By composing jax.vmap with jax.grad and jax.jit, we can get efficient Jacobian matrices, or per-example gradients:

                                                per_example_grads = jax.jit(jax.vmap(jax.grad(loss), in_axes=(None, 0, 0)))

                                                Scaling

                                              To scale your computations across thousands of devices, you can use any composition of these:

                                                * Compiler-based automatic parallelization where you program as if using a single global machine, and the compiler chooses how to shard data and partition computation (with some user-provided constraints);
                                                * Explicit sharding and automatic partitioning where you still have a global view but data shardings are explicit in JAX types, inspectable using jax.typeof;
                                                * Manual per-device programming where you have a per-device view of data and computation, and can communicate with explicit collectives.
                                              Mode      View?       Explicit sharding?  Explicit Collectives?
                                              Auto      Global      ❌                   ❌
                                              Explicit  Global      ✅                   ❌
                                              Manual    Per-device  ✅                   ✅

                                                from jax.sharding import set_mesh, AxisType, PartitionSpec as P
                                                mesh = jax.make_mesh((8,), ('data',), axis_types=(AxisType.Explicit,))
                                                set_mesh(mesh)

                                                # parameters are sharded for FSDP:
                                                for W, b in params:
                                                  print(f' { jax . typeof ( W ) } ')  # f32[512@data,512]
                                                  print(f' { jax . typeof ( b ) } ')  # f32[512]

                                                # shard data for batch parallelism:
                                                inputs, targets = jax.device_put((inputs, targets), P('data'))

                                                # evaluate gradients, automatically parallelized!
                                                gradfun = jax.jit(jax.grad(loss))
                                                param_grads = gradfun(params, (inputs, targets))

                                              See the tutorial and advanced guides for more.

                                                Gotchas and sharp bits

                                              See the Gotchas Notebook.

                                                Installation

                                                Supported platforms

                                                          Linux x86_64  Linux aarch64  Mac aarch64   Windows x86_64  Windows WSL2 x86_64
                                              CPU         yes           yes            yes           yes             yes
                                              NVIDIA GPU  yes           yes            n/a           no              experimental
                                              Google TPU  yes           n/a            n/a           n/a             n/a
                                              AMD GPU     yes           no             n/a           no              experimental
                                              Apple GPU   n/a           no             experimental  n/a             n/a
                                              Intel GPU   experimental  n/a            n/a           no              no


                                                Instructions

                                              Platform         Instructions
                                              CPU              pip install -U jax
                                              NVIDIA GPU       pip install -U "jax[cuda13]"
                                              Google TPU       pip install -U "jax[tpu]"
                                              AMD GPU (Linux)  pip install -U "jax[rocm7-local]"
                                              Intel GPU        Follow Intel's instructions.


                                              See the documentation for information on alternative installation strategies. These include compiling from source, installing with Docker, using other versions of CUDA, a community-supported conda build, and answers to some frequently-asked questions.

                                                Citing JAX

                                              To cite this repository:

                                                @software{jax2018github, author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Yash Katariya and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang}, title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs}, url = {http://github.com/jax-ml/jax}, version = {0.3.13}, year = {2018}, }

                                              In the above bibtex entry, names are in alphabetical order, the version number is intended to be that from jax/version.py, and the year corresponds to the project's open-source release.

                                              A nascent version of JAX, supporting only automatic differentiation and compilation to XLA, was described in a paper that appeared at SysML 2018. We're currently working on covering JAX's ideas and capabilities in a more comprehensive and up-to-date paper.

                                                Reference documentation

                                              For details about the JAX API, see the reference documentation.

                                              For getting started as a JAX developer, see the developer documentation.

                                        About

                                        Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

                                          docs.jax.dev

                                        Topics

                                          jax

                                        Resources

                                          Readme
                                          Apache-2.0 license

                                        Contributing

                                          Contributing
                                          Cite this repository
                                          Activity
                                          Custom properties

                                        Stars

                                          36.1k stars

                                        Watchers

                                          328 watching

                                        Forks

                                          3.7k forks
                                          Report repository

                                        Releases

                                        Used by

                                        Contributors

                                        Languages

    Footer

        © 2026 GitHub, Inc.

      Footer navigation

        * Terms
        * Privacy
        * Security
        * Status
        * Community
        * Docs
        * Contact
        * Manage cookies
        * Do not share my personal information
      You can’t perform that action at this time.