52  Practical Implementation: PyTorch and JAX

This chapter translates neural network concepts into production frameworks:

• scikit-learn — MLPClassifier with standard fit/predict interface (Pedregosa et al. 2011).
• PyTorch — Dynamic computational graphs and object-oriented automatic differentiation (Paszke et al. 2019).
• JAX — Pure functional paradigm with explicit PRNG, vectorization, and XLA compilation (Bradbury et al. 2018; Heek et al. 2023; DeepMind et al. 2020).

1 Deep Learning Framework Ecosystem

Framework	Paradigm	Primary Use Case
scikit-learn	Object-Oriented High-Level	Small tabular data baselines
PyTorch	Imperative / Dynamic Graph	General deep learning & research standard
JAX	Pure Functional / XLA JIT	High-order differentiation & TPU acceleration

from sklearn.neural_network import MLPClassifier
from sklearn.datasets import make_moons

X, y = make_moons(n_samples=400, noise=0.2, random_state=42)
clf = MLPClassifier(hidden_layer_sizes=(16, 16), max_iter=500, random_state=42)
clf.fit(X, y)
print(f"Accuracy: {clf.score(X, y):.3f}")

2 PyTorch Fundamentals

2.1 Tensors and Autograd

\[ \frac{\partial \mathcal{L}}{\partial \theta} \quad \forall \theta \text{ with } \texttt{requires\_grad=True} \tag{1}\]

loss.backward() computes gradients automatically via chain rule backpropagation across dynamic computation graphs.

model.train() — Activates training mode for Dropout/BatchNorm.

optimizer.zero_grad() — Clears accumulated gradients before backward pass.

logits = model(X_train) — Executes forward pass computation graph.

loss.backward() — Computes gradients via autograd reverse graph traversal.

optimizer.step() — Applies weight updates using computed gradients.

Step 1 / 5

Numerically stable Binary Cross Entropy with Logits Loss:

\[ \mathcal{L} = \frac{1}{N} \sum_{i=1}^{N} \left[ \max(z_i, 0) - z_i y_i + \log\left(1 + e^{-|z_i|}\right) \right] \tag{2}\]

3 JAX Fundamentals

JAX relies on pure functions and stateless pytrees.

Mutable object-oriented model encapsulation holding parameters internally.

Stateless pure functions receiving parameters externally as explicit pytree dictionaries.

Single-instance scalar Binary Cross Entropy in JAX:

\[ \mathcal{L}(w) = -\left[ y \log \sigma(wx) + (1-y) \log(1 - \sigma(wx)) \right] \tag{3}\]

import jax
import jax.numpy as jnp
import optax

def loss_fn(params, x, y):
    logits = forward(params, x)
    return optax.sigmoid_binary_cross_entropy(logits, y).mean()

@jax.jit
def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_opt_state, loss

4 Summary

PyTorch provides an intuitive, object-oriented imperative workflow ideal for rapid prototyping and debugging. JAX enforces pure functional programming, enabling high-order automatic differentiation and JIT compilation via XLA (Paszke et al. 2019; Bradbury et al. 2018).

References

Bradbury, James, Roy Frostig, Peter Hawkins, et al. 2018. JAX: Composable Transformations of Python+NumPy Programs. Released. http://github.com/google/jax.

DeepMind, Igor Babuschkin, Kate Baumli, et al. 2020. The DeepMind JAX Ecosystem. Released. http://github.com/google-deepmind.

Heek, Jonathan, Anselm Levskaya, Avital Oliver, et al. 2023. Flax: A Neural Network Library and Ecosystem for JAX. Released. http://github.com/google/flax.

Paszke, Adam, Sam Gross, Francisco Massa, et al. 2019. “PyTorch: An Imperative Style, High-Performance Deep Learning Library.” Advances in Neural Information Processing Systems (NeurIPS) 32.

Pedregosa, Fabian, Gaël Varoquaux, Alexandre Gramfort, et al. 2011. “Scikit-Learn: Machine Learning in Python.” Journal of Machine Learning Research 12: 2825–30. https://jmlr.org/papers/v12/pedregosa11a.html.