What does a JAX model actually learn? Its parameters are the weights and biases that determine a neural network’s computations. Unlike PyTorch or TensorFlow objects, JAX keeps them outside the function that applies the model, making the model state explicit. section 2.2 section 2.4.2
A typical parameter workflow is:
- Initialize weights and biases with JAX’s functional random-number API, passing and splitting PRNG keys explicitly. section 2.4.1
- Pass the parameter structure into a prediction or loss function instead of hiding it in object state. section 2.4.2 section 2.9
- Differentiate the loss with respect to the parameters, then return updated values rather than mutating them in place. section 2.6.3
- Store the collection as a nested Python data structure called a pytree, which can be saved, restored, and composed with transformations. section 2.11 section 2.9
A parameter-driven forward pass
params = init_network_params([28 * 28, 512, 10], random.PRNGKey(0)) preds = predict(params, image) Here, predict receives params explicitly, applies each weight-and-bias pair, and returns logits. The same structure can be supplied to gradient transformations and update functions. section 2.4.1 section 2.4.2
During training, grad(loss)(params, x, y) produces a matching gradient structure. An update subtracts a learning-rate-scaled gradient from every weight and bias; value_and_grad can return the new loss alongside those gradients. Explicit state supports custom updates and works with composable transformations such as jit and vmap. section 2.6.3 section 2.9
Figure 2.1 High-level structure of a JAX project, with data loading, training, and deployment to production
Use the project-structure diagram to place parameters between initialization, training updates, saving, and deployment. section 2.2
