About the project
Contents
About the project#
The JAX project is led by the JAX core team. We develop in the open,and welcome open-source contributions from across the community. Wefrequently see contributions fromGoogleDeepMind, Alphabet more broadly,NVIDIA,and elsewhere.
At the heart of the project is theJAXcore library, which focuses on thefundamentals of machine learning and numerical computing, at scale.
Whendeveloping the core, we want to maintain agilityand a focused scope, so we lean heavily on a surroundingmodulartechnology stack. First, we design thejax moduleto becomposableandextensible, sothat a wide variety of domain-specific libraries can thrive outside ofit in a decentralized manner. Second, we lean heavily on a modularbackend stack (compiler and runtime) to target differentaccelerators. Whether you arewriting a new domain-specific librarybuilt with JAX, or looking tosupportnew hardware, you can oftencontribute these withminimal to no modifications to the JAX corecodebase.
Many of JAX’s core contributors have roots in open-source software andin research, in fields spanning computer science and the naturalsciences. We strive to continuously enable the cutting edge of machinelearning and numerical computing—across all compute platforms andaccelerators—and to discover the truths of array programming atscale.
Open development#
JAX’s day-to-day development takes place in the open on GitHub, usingpull requests, the issue tracker, discussions, andJAX EnhancementProposals(JEPs). Readingand participating in these is a good way to get involved. We alsomaintaindevelopernotesthat cover JAX’s internal design.
The JAX core team determines whether to accept changes andenhancements. Maintaining a simple decision-making structure currentlyhelps us develop at the speed of the research frontier. Opendevelopment is a core value of ours, and we may adapt to a moreintricate decision structure over time (e.g. with designated areaowners) if/when it becomes useful to do so.
For more seecontributing toJAX.
A modular stack#
To enable (a) a growing community of users across numerical domains,and (b) an advancing hardware landscape, we lean heavily onmodularity.
Libraries built on JAX#
While the JAX core library focuses on the fundamentals, we want toencourage domain-specific libraries and tools to be built on top ofJAX. Indeed,manylibraries haveemerged around JAX to offer higher-level features and extensions.
How do we encourage such decentralized development? We guide it withseveral technical choices. First, JAX’s main API focuses on basicbuilding blocks (e.g. numerical primitives, NumPy operations, arrays,and transformations), encouraging auxiliary libraries to developutilities as needed for their domain. In addition, JAX exposes ahandful of more advanced APIs forcustomizationandextensibility. Librariescanlean on theseAPIs inorder to use JAX as an internal means of implementation, to integratemore with its transformations like autodiff, and more.
Projects across the JAX ecosystem are developed in a distributed andoften open fashion. They are not governed by the JAX core team, eventhough sometimes team members contribute to them or maintain contactwith their developers.
A pluggable backend#
We want JAX to run on CPUs, GPUs, TPUs, and other hardware platformsas they emerge. To encourage unhindered support of JAX on newplatforms, the JAX core emphasizes modularity in its backend too.
To manage hardware devices and memory, and for compilation to suchdevices, JAX calls out to the openXLAcompiler and thePJRTruntime. Bothof these are projects external to the JAX core, governed andmaintained by OpenXLA (again, with frequent contributions from anddiscussion with the JAX core developers).
XLA aims for interoperability across accelerators (e.g. by ingestingStableHLO as input) and PJRT offersextensibility through a plug-in device API. Adding support for newdevices is done by implementing a backend lowering for XLA, andimplementing a plug-in device API defined by PJRT. If you’re lookingto contribute to compilation, or to supporting new hardware, weencourage you to contribute at the XLA and PJRT layers.
These open system components allow third parties to support JAX on newaccelerator platforms,without requiring changes in the JAXcore. There are several plug-ins in development today. For example, ateam at Apple is working on a PJRT plug-in to getJAX running onApple Metal.