$ show jax-ml/jax
JAX is a high-performance numerical computing library developed by Google Research. It provides a NumPy-compatible API (jax.numpy), automatic differentiation (jax.grad), just-in-time compilation via XLA (jax.jit), automatic vectorization (jax.vmap), and large-scale parallelization (jax.pmap, jax.shard_map). The library is designed for machine learning research and scientific computing, supporting CPU, GPU (NVIDIA, AMD, Intel), and Google TPU hardware. The system is a hybrid Python/C++ library. The Python layer (jax/) handles function tracing into a JAX intermediate representation (JAXpr), applies transformations (autodiff, batching, sharding), and lowers the computation to MLIR/StableHLO. The C++ backend (jaxlib/) compiles this representation via XLA and manages execution on target hardware. It includes a sophisticated plugin system for hardware backends and a custom kernel DSL (Pallas/Mosaic) for writing high-performance fused operations. While primarily an application/library, JAX has strong research roots. It is explicitly described as a research project, and its experimental/ directory contains rapidly evolving features like custom kernel authoring (Pallas), TensorFlow interop (jax2tf), and sparse array support. The codebase is a blend of production-grade engineering (robust CI, deprecation policies, extensive testing) and cutting-edge research (new autodiff algorithms, novel sharding strategies, experimental hardware support). The repository contains several high-quality, isolated components that are reusable in other Python/C++ projects. These include a file-based LRU cache, a weak-reference LRU cache, thread-safe reentrant hash maps, a lazy module loader, a deprecation framework, and a traceback filtering utility. The core autodiff and JIT compilation patterns are the primary value, but these supporting utilities demonstrate excellent software craftsmanship.
The LLM identified 10 files worth highlighting for reuse:
A file-based LRU cache implementation with atime tracking and thread-safe locking. It provides a decorator that persists cached results to disk, making it ideal for memoizing expensive computations across process restarts. The implementation is clean, well-documented, and handles edge cases like concurrent access and cache eviction robustly.
A high-performance C++ implementation of an LRU cache that uses weak references for its values. This allows cached objects to be garbage collected when no longer referenced elsewhere, preventing memory leaks in long-running processes. The Python bindings are cleanly exposed via pyi stubs.
Thread-safe and reentrant hash map and hash set implementations. These data structures are safe to use inside callbacks or signal handlers, a notoriously difficult problem. The C++ code is meticulously crafted, with extensive testing.
A sophisticated deprecation framework that allows registering deprecations with unique IDs, controlling whether they warn or error, and accelerating the deprecation schedule. It provides a clean, centralized way to manage API churn in a large codebase.
A utility for filtering and formatting tracebacks. It supports defining API boundaries to hide internal frames, integrates with IPython, and provides configurable modes. This significantly improves the developer experience in a complex tracing/JIT environment.
A lazy module loader that defers the import of submodules until they are accessed. This is used extensively in jax.scipy to avoid importing the entire SciPy ecosystem on startup. The implementation is clean and handles submodule attributes correctly.
Implements PartitionSpec (the P object), a core concept in JAX's sharding system. It represents how an array is partitioned across devices. The implementation includes interning for canonicalization, which is a smart optimization for a fundamental data structure.
A set of C macros (FFI_RETURN_IF_ERROR, FFI_ASSIGN_OR_RETURN, CheckShape, AllocateWorkspace) that standardize error handling and boilerplate in XLA FFI handlers. This is a classic example of using C macros to enforce a consistent pattern across a large codebase of handwritten kernels.
An extensible cost estimation framework for Pallas kernels. It uses a registration pattern (register_cost_rule) where different primitives can declare their computational cost (FLOPs, memory bandwidth). This allows the system to estimate the performance of fused kernels without running them.
A simple, clean implementation of a hashable, immutable dictionary (FrozenDict). It is used as a key in caches and as a component of pytrees. Its simplicity and correctness make it a perfect reusable building block.