Chapter 3: Introduction to TensorFlow, PyTorch, JAX, and Keras

Introduction

This chapter introduction previews TensorFlow, PyTorch, and JAX alongside Keras, and points toward practical deep learning applications.

3.1 A brief history of deep learning frameworks

This section traces how deep learning frameworks developed and explains why their shared capabilities and today’s ecosystem matter to practitioners.

• Three framework capabilities: Explain the three capabilities that modern deep learning frameworks combine: automatic differentiation, tensor computation on CPUs and GPUs, and computation distributed across devices or computers.
  Why it matters: These shared capabilities explain why practitioners use frameworks instead of writing low-level training code from scratch.
• Early framework turning points: Trace why practical deep learning frameworks emerged relatively recently, from hand-written gradients and difficult GPU programming to CUDA and Theano.
  Why it matters: The history shows why having these capabilities together in a reusable tool was a major change.
• Modern framework lineage: Describe the main relationships in the section’s framework history: Keras began on Theano, TensorFlow added large-scale distributed computation, PyTorch drew on Chainer and NumPy-Autograd, and JAX offered autodifferentiation with XLA.
  Why it matters: These relationships make the similarities and differences among the four current frameworks easier to place in context.
• Ecosystem stability outlook: Explain the author’s outlook that Python and multiple frameworks are likely to persist, while Keras backends and compatibility with existing frameworks can help workflows survive new frameworks or chips.
  Why it matters: This outlook explains why the section recommends learning a little about each framework without assuming today’s tools will be the last ones.

Key terms: automatic differentiation, tensor computations, distributed computation, CPUs, GPUs, CUDA, Theano, autodifferentiation, GPU computation, Keras

3.2 How these frameworks relate to each other

This section distinguishes Keras’s high-level role from the lower-level roles of TensorFlow, PyTorch, and JAX, and explains why Keras needs a backend.

• High-level and low-level roles: Distinguish Keras’s role from those of TensorFlow, PyTorch, and JAX. Describe the training concepts associated with each level.
  Why it matters: This distinction helps you recognize what each framework contributes to building and training a model.
• Keras backend requirements: Explain why Keras needs a backend and why a NumPy backend cannot support training.
  Why it matters: The backend determines whether a Keras workflow has the gradient support needed for training.

Key terms: high-level framework, low-level APIs, tensors, backpropagation, layers, loss function, backend engine, Keras backends, NumPy, gradients

3.3 Introduction to TensorFlow

Introduces TensorFlow as a widely used machine learning framework and a broader ecosystem spanning research and production.

• TensorFlow platform and ecosystem: Explain that TensorFlow is a Python-based open source framework developed primarily by Google, but is more than a single library. Describe how its ecosystem supports uses from research to production, giving examples of components and their purposes.
  Why it matters: Thinking of TensorFlow as a platform explains why its name comes up in contexts beyond building a model.

Key terms: TensorFlow, TFX, TF-Serving, TF Optimization Toolkit, TFLite, MediaPipe

3.3.1 First steps with TensorFlow

This section introduces TensorFlow tensor creation, mutable variables, automatic gradients, and function compilation.

• Tensor creation methods: Describe how to create TensorFlow tensors with initial values using tf.ones, tf.zeros, tf.constant, or tf.random methods.
  Why it matters: These are the starting points for creating values used in TensorFlow computations.
• Mutable model state: Distinguish non-assignable TensorFlow tensors from tf.Variable, and describe how a variable's value can be updated.
  Why it matters: Variables make model-state updates possible even though ordinary TensorFlow tensors cannot be assigned to.
• GradientTape gradient retrieval: Explain how a GradientTape retrieves gradients of a differentiable computation with respect to its inputs, including its use for model loss and weights.
  Why it matters: Gradients provide the information used to update model weights during training.
• Gradient tracking choices: Explain which inputs GradientTape tracks by default, how to track a constant tensor, and why tracking is selective.
  Why it matters: Knowing what the tape watches helps explain why a requested gradient may require an explicit step.
• Second-order gradients: Describe what a second-order gradient is and how nested gradient tapes can compute one.
  Why it matters: The position example shows that GradientTape can differentiate a computed gradient, not just an original expression.
• TensorFlow function compilation: Explain how compiling a TensorFlow function differs from eager execution, why and when to use it, and how graph mode and XLA compare.
  Why it matters: This helps you choose between straightforward debugging and faster repeated execution.

Key terms: tf.ones, tf.zeros, tf.constant, tf.random, TensorFlow tensors, tf.Variable, assign, assign_add, assign_sub, GradientTape

3.3.2 An end-to-end example: A linear classifier in pure TensorFlow

The section builds a linear classifier for two separable classes of 2D points and shows how its learned predictions define a separating line.

• Synthetic points and labels: Describe how the example organizes two classes of 2D points as inputs and assigns their target labels.
  Why it matters: The input and target shapes make clear what the classifier receives and what it is trained to predict.
• Affine predictions and squared loss: Explain how the classifier computes a prediction from a 2D point and what its training seeks to minimize.
  Why it matters: This connects the model’s adjustable parameters to both its predictions and its training objective.
• Whole dataset training steps: Distinguish this example’s batch training from mini-batch training and explain the stated tradeoff.
  Why it matters: It explains why the example uses relatively few, more expensive training steps.
• Threshold and separating line: Trace how a numerical prediction becomes a class label and why the classifier’s boundary is a line.
  Why it matters: It shows how learned numbers become a geometric decision boundary.

Key terms: negative_samples, positive_samples, inputs, targets, linear classifier, affine transformation, W, b, loss function, batch training

3.3.3 What makes the TensorFlow approach unique

This section compares TensorFlow’s strengths and limitations to help readers judge when it may be a good framework choice.

• TensorFlow strengths and use cases: Explain why TensorFlow can be a strong choice for speed, varied data, preprocessing, and production deployment.
  Why it matters: These strengths help explain why TensorFlow may suit a project beyond simply running model computations.
• TensorFlow adoption drawbacks: Describe the limitations that may make TensorFlow harder to use or less suitable for some pretrained-model work.
  Why it matters: A framework’s breadth does not guarantee that it will be the easiest choice for a particular project.

Key terms: graph mode, XLA compilation, string tensors, ragged tensors, tf.data, production deployment, sprawling API, NumPy API, Hugging Face, pretrained models

3.4 Introduction to PyTorch

Introduces PyTorch’s place in the machine-learning ecosystem and compares its API with TensorFlow, JAX, and Keras.

• PyTorch API and Keras compatibility: Explain that PyTorch provides higher-level components such as layers and optimizers, and that these components are compatible with Keras workflows when Keras uses the PyTorch backend.
  Why it matters: This comparison helps you place PyTorch alongside the other frameworks introduced in the chapter.
• PyTorch origins and ecosystem: Describe PyTorch as an open-source, Python-based framework developed primarily by Meta, used extensively in machine-learning research, and surrounded by related packages and platforms.
  Why it matters: This context explains why PyTorch and its surrounding tools appear throughout machine-learning work.

Key terms: PyTorch API, layers, optimizers, Keras workflows, PyTorch backend, PyTorch, Meta, torchvision, torchaudio, Hugging Face

3.4.1 First steps with PyTorch

This section introduces PyTorch tensors and parameters, then explains how backward() computes gradients and why those gradients sometimes need resetting.

• PyTorch tensor conventions: Describe how PyTorch is imported and how its tensor creation and assignment conventions compare with NumPy and TensorFlow.
  Why it matters: These conventions help you read and write the PyTorch examples without assuming its API exactly matches NumPy or TensorFlow.
• Parameters as trainable state: Explain why PyTorch provides Parameter in addition to regular tensors for model state.
  Why it matters: Parameter distinguishes the values a model trains from tensors used for other purposes.
• Backward gradient computation: Trace what happens when backward() is called on a PyTorch tensor after a computation.
  Why it matters: This is how PyTorch exposes the gradients needed for training without an explicit gradient-tape API.
• Gradient accumulation and reset: Explain what repeated backward() calls do to .grad and state how to reset it as shown here.
  Why it matters: Without accounting for accumulation, a stored gradient may include earlier backward() calls.

Key terms: torch, tensor, normal, torch.rand, torch.nn.parameter.Parameter, torch.Tensor, trainable state, computation graph, tape, tensor.backward()

3.4.2 An end-to-end example: A linear classifier in pure PyTorch

This section builds a linear classifier in PyTorch, first with manual gradient updates and then with a Module and optimizer, before introducing optional compilation.

• Manual gradient training step: Trace how a training step obtains gradients of the loss with respect to W and b, updates them without recording gradients, and clears their stored gradients.
  Why it matters: These steps explain both how the classifier learns and why leaving old gradients in place breaks training.
• Module state and forward computation: Describe how a torch.nn.Module subclass packages model parameters and the forward computation, and how an instance is called on inputs.
  Why it matters: The Module groups the classifier’s variables and computation so they can be used together in training.
• Optimizer-based parameter updates: Explain how the Module’s parameters are supplied to a PyTorch optimizer and how the simplified training step replaces manual updates and gradient resets.
  Why it matters: This is the shorter training-step pattern the section recommends, with gradient clearing as an essential part.
• Optional PyTorch compilation: Describe what torch.compile() is intended to do and why PyTorch code may still run eagerly.
  Why it matters: Compilation offers a possible performance improvement, not a guaranteed requirement for using PyTorch.

Key terms: loss.backward(), tensor.grad, torch.no_grad(), W, b, torch.nn.Module, Parameters, __init__(), forward(), __call__()

3.4.3 What makes the PyTorch approach unique

This section weighs PyTorch’s debugging and model-availability strengths against its API inconsistencies and performance drawbacks.

• Eager execution and debugging: Explain why PyTorch’s emphasis on eager execution makes debugging easier, and distinguish that emphasis from TensorFlow and JAX.
  Why it matters: This explains a practical reason to choose PyTorch when ease of debugging matters.
• Hugging Face model availability: Explain how Hugging Face support contributes to PyTorch adoption.
  Why it matters: Ready access to pretrained models is a major practical reason people use PyTorch.
• API naming inconsistencies: Describe why PyTorch’s API can be frustrating to learn, especially for someone coming from NumPy.
  Why it matters: Inconsistent conventions can make familiar-looking code harder to learn and use.
• Eager execution performance tradeoff: Explain the performance downside the section associates with PyTorch’s focus on eager execution and the limitation it gives for torch.compile().
  Why it matters: The ease of eager debugging comes with a performance tradeoff that compilation may not fully resolve.

Key terms: eagerly, compilation, debugging, Hugging Face, first-class support, pretrained model, PyTorch API, NumPy, axis, dim

3.5 Introduction to JAX

Introduces JAX and explains how its stateless approach supports large-scale machine learning computation.

• Stateless computation and scalability: Explain what stateless computation means in JAX and why it helps JAX scale.
  Why it matters: This connects JAX’s programming approach to its usefulness for very large machine learning problems.
• JAX purpose and uptake: Describe what JAX is and why the section says it gained traction in research.
  Why it matters: This places JAX’s scaling strengths in the context of its research use.

Key terms: stateless, automatic parallelization, distributed computation, synchronization, JAX, differentiable computation, TPUs

3.5.1 First steps with JAX

This subsection previews the JAX concepts that will be covered next.

3.5.2 Tensors in JAX

This section introduces JAX tensors through its NumPy-compatible array API and the `jax.numpy` namespace.

• JAX NumPy API: Explain that JAX uses the NumPy API rather than a separate, slightly different numerical API, and identify `jax.numpy` and its common alias `jnp`.
  Why it matters: Recognizing `jnp` helps you read JAX array code using familiar NumPy conventions.

Key terms: NumPy API, jax.numpy, jnp

3.5.3 Random number generation in JAX

This section explains how JAX’s stateless design affects random numbers and array updates, then shows how to compute gradients and compile functions.

• Explicit random seed keys: Explain why JAX random operations require explicit keys, how to obtain new keys, and why this supports reproducible, vectorizable code.
  Why it matters: Explicit keys make random behavior repeatable without relying on hidden state.
• New arrays for updates: Describe how JAX updates an array element without changing the original array in place.
  Why it matters: Treating updates as new values helps keep changes to JAX data explicit.
• Gradient functions from losses: Trace the basic steps for computing gradients with jax.grad(), including the required loss output and the argument differentiated.
  Why it matters: This is the basic pattern for obtaining gradients needed in JAX training code.
• Structured gradient state: Describe how the first argument of a loss function can hold multiple values for which JAX computes gradients.
  Why it matters: Grouping state lets the same gradient pattern handle losses involving multiple variables.
• Loss and auxiliary returns: Distinguish jax.value_and_grad() from jax.grad(), and describe how has_aux handles an extra output from the loss function.
  Why it matters: These options avoid separately recomputing the loss and let training code retain other useful outputs.
• Compilation of stateless functions: Explain what @jax.jit does and the stateless-function constraint on its use.
  Why it matters: The speed benefit of JIT compilation comes with a requirement to make updated values explicit.

Key terms: PRNG, key, jax.random.split(), statelessness, tensor assignment, at()/set() API, in-place modification, jax.grad(), compute_loss(), grad_fn

3.5.4 An end-to-end example: A linear classifier in pure JAX

This section shows how to build a JAX linear-classifier training loop using stateless functions and automatic gradient computation.

• Stateless training state flow: Trace how W and b move through the JAX training loop: they are passed into stateless functions, updated values are returned by training_step(), and the loop repeatedly uses those returned values. Explain why the fixed learning rate is not passed as state in this example.
  Why it matters: Explicitly passing updated state is what makes the JAX training step work without retaining model state inside the function.
• Loss and gradient computation: Explain how the example packages scalar loss computation in compute_loss() and uses jax.value_and_grad() to obtain both the loss and gradients with respect to the tensors in state.
  Why it matters: This connects the classifier’s loss calculation to the gradients needed to update its parameters.

Key terms: stateless, state, W and b, training_step(), learning_rate, compute_loss(), scalar, jax.value_and_grad(), gradients

3.5.5 What makes the JAX approach unique

This section explains JAX’s functional, stateless approach and weighs its performance advantages against its development costs.

• Functional stateless JAX philosophy: Explain that JAX’s distinguishing approach is functional and stateless, and that this approach enables fast compiled code and scaling to large models and many devices.
  Why it matters: It connects JAX’s distinctive programming style to the benefits that motivate using it.
• JAX performance and platform advantages: Describe the practical advantages the section attributes to JAX: speed for most models, a NumPy-consistent numerical API, and suitability for TPU training.
  Why it matters: These advantages help explain when JAX may be an appealing framework choice.
• JAX development tradeoffs: Explain why JAX may be harder to work with: metaprogramming and compilation complicate debugging relative to pure eager execution, while low-level training loops can be more verbose than in TensorFlow or PyTorch.
  Why it matters: It gives a balanced view of JAX rather than treating speed as the only consideration.

Key terms: functional, stateless philosophy, compile, NumPy, TPUs, XLA, metaprogramming, compilation, pure eager execution, low-level training loops

3.6 Introduction to Keras

This section introduces Keras as a Python deep learning API and explains why it offers multiple ways to build and train models.

• Keras API purpose: Describe Keras as a Python deep learning API that provides a convenient way to define and train deep learning models.
  Why it matters: This establishes what Keras is for before you encounter its model-building tools.
• Range of Keras workflows: Explain why Keras supports multiple ways to build and train models, and describe the tradeoff those ways represent.
  Why it matters: The range of workflows explains why different Keras users may build and train models differently.

Key terms: Keras, deep learning API, workflows, high-level, low-level, usability, flexibility

3.6.1 First steps with Keras

This section explains how Keras uses interchangeable backend frameworks and how to choose one before importing Keras.

• Keras backend portability: Explain what a Keras backend does and why switching backends can be useful for the same model.
  Why it matters: You can choose the framework ecosystem that suits each stage of working with a Keras model.
• Backend selection before import: Describe Keras’s default backend and how to select another one before importing Keras.
  Why it matters: Choosing the backend at setup time ensures Keras runs on the framework you intend to use.

Key terms: Keras, backend frameworks, JAX, TensorFlow, PyTorch, default backend, KERAS_BACKEND, ~/.keras/keras.json, import keras

3.6.2 Layers: The building blocks of deep learning

This section explains Keras layers as stateful data-processing building blocks and shows how they infer input shapes and create weights when first called.

• Layer state and computation: Explain what a layer takes in and produces, and how a Keras Layer combines learned weights with a forward-pass computation.
  Why it matters: This is the basic unit you assemble into a Keras model.
• Layers for tensor formats: Describe how the section matches dense, recurrent or 1D convolution, and 2D convolution layers to different data and tensor formats.
  Why it matters: The data's structure helps determine which kind of layer is appropriate.
• Compatible layer shapes: Explain what makes adjacent layers compatible and why automatic input-shape inference simplifies assembling Keras models.
  Why it matters: Automatic shape inference avoids specifying each layer's input size by hand while preserving shape compatibility.
• First-call weight creation: Trace the roles of build(), call(), and __call__() when a custom Keras layer is used for the first time.
  Why it matters: Keeping weight creation separate from the forward pass lets a custom layer adapt to its first input.

Key terms: layer, weights, forward pass, dense layers, Dense, recurrent layers, LSTM, Conv1D, Conv2D, layer compatibility

3.6.3 From layers to models

This section explains how Keras models organize layers into networks and why a model’s topology determines what it can learn.

• Models as Layer Graphs: Describe a Keras model as a graph of layers, and distinguish the simple single-input, single-output stack of a Sequential model from other possible network topologies.
  Why it matters: It helps you recognize that the layer stacks you have used so far are only one kind of model.
• Architecture and Learnable Possibilities: Explain how choosing a network topology limits the tensor operations and weight values a model can use to map inputs to outputs. Describe how that choice encodes assumptions about the problem, using the single Dense layer example.
  Why it matters: Architecture matters because adjusting weights cannot remove the assumptions imposed by the model’s structure.

Key terms: Model class, Sequential models, graph of layers, network topology, hypothesis space, weight tensors, architecture, linearly separable

3.6.4 The “compile” step: Configuring the learning process

This section explains the choices needed to configure Keras training and how to supply them to compile().

• Training configuration choices: Distinguish the roles of the loss, optimizer, and metrics when configuring training.
  Why it matters: These distinctions explain what training changes and what its reported measures tell you.
• Compile argument forms: Describe how compile() configures training and distinguish passing its settings as strings from passing object instances.
  Why it matters: It lets you choose between concise built-in defaults and explicitly configured training components.

Key terms: loss function, objective function, optimizer, metrics, stochastic gradient descent (SGD), compile(), loss, object instances, learning_rate

3.6.5 Picking a loss function

This section explains why a loss must reflect the intended task and gives standard loss choices for common classification problems.

• Loss and task alignment: Explain why the loss function must match the intended outcome: a network minimizes the specified loss, even when doing so produces unwanted results.
  Why it matters: A model can optimize exactly what you specified without achieving what you intended.
• Standard classification loss choices: State the section’s loss choices for two-class and many-class classification, and distinguish these common cases from novel research problems that may require a custom loss.
  Why it matters: For familiar tasks, established guidelines give you a starting point for choosing a suitable loss.

Key terms: loss function, objective, shortcut, unintended side effects, binary crossentropy, categorical crossentropy, two-class classification, many-class classification, loss functions

3.6.6 Understanding the fit method

This section explains what Keras fit() does, its main training arguments, and the per-epoch results it returns.

• Fit training inputs and schedule: Explain that fit() runs the training loop after compile(), using training inputs and targets, a number of epochs, and a batch size.
  Why it matters: These arguments determine what data fit() trains on and how its training loop proceeds.
• Returned training history: Describe the History object returned by fit() and what its history field records.
  Why it matters: The returned history lets you inspect how reported loss and metrics changed across epochs.

Key terms: fit, inputs and targets, epochs, batch size, mini-batch gradient descent, History object, history, loss, metrics, per-epoch values

3.6.7 Monitoring loss and metrics on validation data

The section explains why validation data must be kept separate from training data and how Keras reports loss and metrics on it during or after training.

• Generalization Beyond Training Data: Explain why good performance on training data does not establish that a model will perform well on unseen data.
  Why it matters: Training performance alone cannot tell you whether the model will be useful on new examples.
• Separate Validation Data Monitoring: Describe how to monitor performance on new data with validation data in fit(), and why that data must remain separate from training data.
  Why it matters: A separate validation set lets you monitor whether training is producing something useful on new data.
• Post Training Evaluate Results: Describe how evaluate() computes validation loss and metrics after training and what it returns.
  Why it matters: It gives you a way to obtain validation results after fitting rather than only during it.

Key terms: training data, general, memorizing, validation data, validation_data, fit(), validation loss, training loss, evaluate(), validation metrics

3.6.8 Inference: Using a model after training

This section explains how to use a trained Keras model for inference and why predict() is preferable to calling the model directly on a large input set.

• Inference with batched predictions: Explain that inference means using a trained model to predict on new data, and describe why Keras predict() is preferable to calling the model directly when there are many inputs.
  Why it matters: Batching lets you obtain predictions without requiring all the inputs to fit in GPU memory at once.

Key terms: inference, __call__(), predict(), small batches, NumPy array, TensorFlow Dataset objects

3.7 Summary