What you will learn in Deep Learning with JAX
Chapter by chapter
Concepts explained in this book
- just-in-time compilationCompilation performed at runtime so code can be specialized and optimized for actual inputs and hardware.
- parallel mapThe JAX mapping transformation that parallelizes a computation across multiple local devices.
- JAXA functional-first numerical computing library for machine learning and scientific computing in which state is made explicit and transformations like autodiff and compilation are central.
- jax.gradA JAX transformation that returns a derivative-computing version of a function, typically with respect to its first argument by default.
- collective operationA parallel-programming operation that communicates or combines values across mapped devices or processes.
- automatic differentiationA technique for computing exact derivatives of code-expressed computations by systematically applying differentiation rules to the program structure.
- OptaxA JAX-based library of gradient transformations and optimizer-building blocks used to define and combine training updates.
- in_axesThe vmap() argument that specifies, for each function argument, which axis to map over or whether to leave it unmapped.
- language modelA large neural language model used for language-modeling tasks such as generation, training, fine-tuning, and serving.
- keyA JAX pseudorandom-number-generator key passed explicitly to random functions for functional and reproducible randomness.
- pytreeA recursive tree-like structure built from Python container nodes and leaf values, used in JAX to organize nested function arguments and results.
- gradient descentAn optimization step that updates model parameters using gradients, typically in the opposite direction of the loss gradient.
- forward-mode autodiffAn automatic differentiation approach that propagates derivative information backward from outputs to inputs.
- splitThe act of dividing one random-number-generator key into multiple independent keys for separate random draws.
- parameterThe model parameters stored in a training state and updated by an optimizer during training.
- side effectA program behavior that changes state outside the function’s explicit return value, such as mutating hidden variables or global state.
- jax.device_put()A JAX function that transfers data to a specified device and can commit it to that device, returning without necessarily waiting for completion.
- composable function transformationA higher-order operation that takes a function and returns a transformed function, such as a compiled, differentiated, or vectorized version.
- functional, stateless philosophyA programming style centered on functions as values and minimal reliance on mutable state.
- prompt engineeringThe practice of shaping input prompts to get better outputs from a pretrained language model without changing its weights.
- SPMDA parallel programming model in which the same program runs simultaneously on multiple devices while each device processes different data.
- apply()The Linen method that runs a model forward pass using provided variables and inputs, optionally updating mutable state.
- init()The Linen method that initializes a Module's variables, including parameters, from dummy inputs and a PRNG key.
- gradientThe derivative of a function with respect to its inputs, used here in the optimization sense for training neural networks.
- static argumentA function argument whose concrete value is fixed during tracing and whose changes trigger recompilation.
- out_axesThe parameter that specifies which output axis should hold the mapped dimension in a transformed function’s result.
- XLAA compiler and runtime for linear algebra and array programs that optimizes computations such as those produced by JAX.
- LinenFlax's neural-network API for defining models as Modules in a functional JAX style.
- parallelismThe execution of computation concurrently across multiple devices or cores.
- GPT-J-6BAn open-source GPT-like autoregressive language model with 6B parameters trained on the Pile and used here as the main pretrained example.
- ahead-of-time compilationCompilation performed before execution time so the resulting code can run without a warm-up JIT step.
- jax.numpyJAX's NumPy-compatible array API module for creating and manipulating arrays.
- axis nameLabels assigned to mapped axes in JAX so collective operations can refer to the device or batch dimension they act over.
- JaxprJAX's simple intermediate representation of a traced function, used as the input to later transformations and compilation.
- TrainStateA Flax dataclass-like container that bundles model parameters, optimizer state, step count, and an apply function for training.
- device meshAn arrangement of multiple devices in a grid used to assign different tensor shards to different devices for parallel computation.
- auto classA Transformers abstraction that selects and instantiates the appropriate tokenizer or model class from a pretrained model name and backend.
- FlaxAutoModelForCausalLMA Hugging Face Flax auto class for loading pretrained causal language models.
- argnumsThe grad-related parameter that selects which positional argument or arguments the derivative should be taken with respect to.
- has_auxA grad-related parameter indicating that a transformed function returns auxiliary data in addition to the differentiated output.
- noiseUnstructured pseudo-random perturbation added to training examples, here used as an image augmentation.
- Common Loop UtilsA library that provides common utilities for machine learning training loops and functional metric computation.
- jax.lax.stop_gradient()A low-level JAX operation that passes values through unchanged while blocking gradient propagation through them.
- data parallelA distributed training strategy that splits a batch across devices, replicates parameters, and aggregates gradients before updating.
- jacfwd()A JAX function that computes a Jacobian by applying forward-mode differentiation repeatedly to input directions.
- JacobianThe matrix of all first-order partial derivatives of a vector-valued function with respect to multiple inputs.
- jax.devices()A JAX function that returns the full set of devices available across processes.
- embeddingA rank-1 multidimensional array with one axis, used here for model outputs such as class probabilities.
- code compilationThe process of converting code into optimized executable form.
- jax.process_index()A JAX function that returns the integer index of the current process in a multi-process program.
- FlaxA high-level neural-network library built on JAX that provides a functional core and module-based APIs.
- JVPThe product of a function's Jacobian with an input vector, computed here via forward-mode autodiff without forming the full Jacobian.
- data augmentationThe practice of creating modified training examples to increase diversity and improve model robustness.
- device placementThe policy of deciding which device or set of devices owns an array and where its computations execute in JAX.
- jax.lax.psumA JAX collective reduction that sums values across a mapped device or process axis.
- NumPyA Python library for array-based numerical computing.
- @nn.compactA Linen annotation that lets a model's forward computation and submodule definitions be written directly inside a single method.
- keyAn explicit RNG state value passed to a randomized JAX function and typically split for subsequent uses.
- committed dataJAX array data that is not explicitly bound to a device and can be moved or resharded automatically.
- jax.device_get()A JAX function that transfers data from a device back to the Python host process as a NumPy ndarray.
- jax.make_jaxpr()A JAX transformation that returns the Jaxpr corresponding to a Python function call signature.
- allreduceA collective communication pattern that reduces values across devices and makes the result available to all participants.
- JAX2TFA converter that translates JAX models or computations to TensorFlow format for interoperability.
- mirrorAn array or image transformation that reverses the order of elements along a chosen dimension.
- OrbaxA JAX-oriented library for checkpointing, serialization, and exporting models to SavedModel.
- random augmentationA data-augmentation step that randomly chooses one transformation to apply to an input image at runtime.
- activation functionA nonlinear function applied inside a neural network to let stacked layers represent more than a single linear transformation.
- jax.value_and_gradThe JAX function that returns both a function value and its gradient in one transformation.
- Gaussian noiseRandom noise with a Gaussian distribution added to data, here used to simulate image corruption.
- batch normalizationA neural-network normalization layer that normalizes activations using batch statistics during training and stored running statistics during inference.
- jax.random.normal()A JAX random-number function that draws samples from a normal distribution using an explicit key and output shape.
- immutable collectionThe property of values not being modified in place after creation, as emphasized for JAX arrays here.
- pseudo-random number generatorA deterministic algorithm that produces numbers that appear random from an internal state or seed.
- jax.laxThe lower-level JAX API that provides primitive operations underlying higher-level array libraries such as jax.numpy.
- forward passThe computation that applies a neural network to input data to produce output activations or predictions.
- jax.tree_util.tree_map()A JAX utility that applies a function across corresponding leaves of one or more pytrees and returns a new pytree.
- training datasetThe dataset used to compute gradients and update model parameters during training.
- jax.distributed.initialize()The JAX API call that initializes distributed execution across multiple processes and hosts.
- multihost configurationA distributed JAX setup in which the same program runs across multiple processes, hosts, and devices.
- JAX ONNX RuntimeA toolchain that executes ONNX models with JAX as the backend.
- SaxmlAn experimental serving system for JAX-, Paxml-, and PyTorch-based models used for inference.
- TF2JAXAn experimental library for converting TensorFlow functions or graphs into JAX functions.
- structured control flowControl-flow constructs in JAX that can be compiled from traces, often using lower-level lax APIs rather than ordinary Python runtime branching.
- diffusersA library of pretrained and trainable diffusion-model pipelines for generating images, audio, video, and other outputs.
- diffusion modelA generative model that iteratively denoises noise to synthesize outputs such as images, audio, or 3D structures.
- FlaxStableDiffusionPipelineA Flax-based pipeline class for loading and running Stable Diffusion models with JAX and Flax.
- replicate()A Flax utility that copies parameters or other arrays onto multiple devices for parallel execution.
- shard()A Flax utility that partitions arrays or pytrees across local devices for parallel computation.
- Stable DiffusionA widely used pretrained text-to-image diffusion model.
- directional derivativeThe derivative of a function in a specified input direction, often used to explain Jacobian-vector products.
- primalThe original function values carried alongside derivative information during forward-mode autodiff.
- tangent vectorDerivative values propagated alongside primals in forward-mode autodiff to represent sensitivity to inputs.
- JAX ecosystemThe collection of libraries and modules built around JAX for different machine-learning tasks.
- foundational modelA model that has already been trained on a large corpus and can be reused directly or adapted to a new task.
- bfloat16A 16-bit floating-point format with a float32-like exponent range and reduced significand precision, commonly used for deep learning.
- float16A 16-bit floating-point numeric type used as a lower-precision data type for deep learning and inference.
- float64A 64-bit floating-point numeric type used as a higher-precision array data type in numerical computing.
- jax.experimental.sparseJAX’s experimental module for sparse tensor and sparse matrix operations.
- sparse representationA tensor representation that stores only non-zero values to reduce memory usage.
- type promotion semanticsRules that determine the result dtype of an operation when its operands have different data types.
- lax.condA JAX control-flow primitive that selects one of two branches based on a predicate while remaining traceable and compilable.
- lax.fori_loopA JAX control-flow primitive for looping over a fixed integer range while remaining traceable and compilable.
- cotangentA reverse-mode sensitivity value associated with a variable, representing the derivative of an output with respect to that variable.
- vector-Jacobian productThe product of an incoming vector with the Jacobian of a function, computed to propagate sensitivities efficiently.
- inner productA linear algebra operation that multiplies corresponding elements of two arrays and sums the products.
- lax.switchA JAX control-flow primitive that selects one branch from multiple alternatives based on an integer index.
- Cloud TPUGoogle Cloud's TPU deployment architecture for running virtual machines attached to TPU accelerators.
- TPUA tensor processing unit hardware backend used to accelerate machine-learning computations.
- fold_in()A JAX random utility that combines an existing key with integer data to derive a new pseudo-random key.
- ShapedArrayA traced abstract value that carries shape information but not concrete runtime values, used by JIT tracing in JAX.
- IvyA framework and transpiler that unifies multiple ML backends and can convert code among JAX, TensorFlow, PyTorch, and NumPy.
- Keras 3A multi-backend neural-network library that can run with JAX, TensorFlow, and PyTorch backends.
- locally optimal solutionA point where a function has a value no greater than that of nearby points.
- filteringThe process of transforming an image by applying a kernel or filter to produce a modified output.
- Gaussian blurAn image-smoothing filter that uses Gaussian-shaped weights to reduce noise while blurring detail.
- GPU kernelA compiled low-level executable unit for performing an operation on target hardware, not a convolution kernel.
- naive approachA simple but potentially inefficient way to apply a scalar function to arrays, either by relying on accidental broadcasting or by looping over elements manually.
- SwishA smooth activation function discovered by automatic search and used here as the example nonlinearity in the network.
- analytical differentiationA differentiation method that derives an exact algebraic expression for a derivative before evaluating it.
- AlpaA system for automatically parallelizing and training or serving large-scale neural networks on distributed clusters.
- EasyLMA JAX/Flax library for pretraining, fine-tuning, evaluating, and serving large language models.
- JaxformerA minimal JAX library for training large language models on TPUs using data and model parallelism.
- JaxSeqA JAX-based library built on Transformers for training large language models and supporting several pretrained model families.
- LevanterA framework for training large language models and other foundation models in JAX with emphasis on legibility, scalability, and reproducibility.
- MaxTextA simple, scalable JAX library for training large language models in pure Python/JAX, especially on Google Cloud TPUs.
- PaxmlA JAX-based framework for training very large models across TPU slices or pods with configurable parallelization and experimentation.
- T5XA modular JAX-based framework for high-performance training, evaluation, and inference of language sequence models.
- MNIST datasetA standard handwritten-digit image dataset used for machine learning experiments and evaluation.
- weight initializationThe process of assigning initial values to neural-network weights before training begins.
- sharpening filterA filter designed to increase local contrast and make image details appear crisper.
- image filterAn image-processing operation that transforms pixels to produce a visual effect or enhance the image.
These guides are generated from the text of Manning books by Manning's knowledge graph and link to the sections they draw on. The full books are available on liveBook and at manning.com.