How deeply does Deep Learning with JAX cover apply()?

Explains the call signature and state return, but defers deeper stateful examples to later state-management material.

Apply()

The Linen method that runs a model forward pass using provided variables and inputs, optionally updating mutable state.

Runs a model forward pass with explicit variables and inputs, instead of keeping parameters inside the module. It replaces older predict-style calls and can also return updated mutable state when the model has one. section 11.1.1 section 11.4 section 11.2.1

Typical use here is to swap in apply() at the places that need model outputs during training and evaluation.

Call it on an instantiated Module with initialized variables and input data. section 11.1.1 section 11.4
Use it in the loss function and accuracy calculation instead of predict(). section 11.1.1
Pass mutable collections when layers like BatchNorm need state updates. section 11.2.1

MLP forward pass

An MLP is initialized from a dummy input, then model.apply(params, images) produces logits for the loss function and training update step. section 11.1.1 section 11.4

Model state update

With mutable state, apply_fn receives {'params': params, **model_state}, an input batch, mutable keys, and train=True, then returns output plus updated state. section 11.2.1

How Deep Learning with JAX covers it

liveBook manning.com

Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

TrainStateused by · A Flax dataclass-like container that bundles model parameters, optimizer state, step count, and an apply function for training.
moduleequivalent to · A module is a distinct, bounded software unit that groups related code or functionality and provides an organizational or interaction boundary within a larger system.
Flaxequivalent to · A high-level neural-network library built on top of JAX that provides a functional core and a module-based API.
jax.gradused by · jax.grad is a JAX transformation that creates a function computing derivatives, typically gradients of a scalar-output function with respect to specified inputs.

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.