Jax

# PyTorch and JAX
import torch
from torch.utils.data import TensorDataset, DataLoader

X_raw = torch.randn(1024, 10)
Y_raw = torch.randn(1024, 1)

# simply combines them so indexing is easier
dataset = TensorDataset(X_raw, Y_raw) 
# for batch loading
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# PyTorch

# some already imported above
#import torch
#from torch.utils.data import TensorDataset, DataLoader
import torch.nn as nn
import torch.optim as optim

# stateful model
class MLP(nn.Module):
	def __init__(self):
		super().__init__()
		# note the convention for dimensions: from 10 to 32 for example
		# in math we'd usually have W*x where W is (out_dims, in_dims)
		self.net = nn.Sequential(
			nn.Linear(10, 32),  
			nn.ReLu(),
			nn.Linear(32, 1)
		)
		
	def forward(self, x):
		return self.net(x)
		
model = MLP()
criterion = nn.MSELoss()
# the optimiser holds a reference to the parameters so it can .step() them
optimiser = optim.Adam(model.parameters(), lr=0.01)

# classic training loop
for epoch in range(5):
	for batch_x, batch_y in dataloader:
		# zeroes out the `.grad` attribute stored on each parameter tensor
		optimiser.zero_grad()
		
		preds = model(batch_x)
		loss = criterion(preds, batch_y)
		
		# add gradients to each parameter via chain rule
		loss.backward()
		
		# now uses the gradients to update parameters
		optimiser.step()
		

# Raw JAX

import jax
import jax.nn as nn # has some helpers, but we don't have a loss function
import jax.numpy as jnp
import optax # for our optimiser

#
# Setup definitions
#

# notice how randomness needs to be used throughout (it's not part of 
# hidden state somewhere)

# params are just a pytree (nested dict) — no object holds them
# we also need to handle initialisation (PyTorch does this for us)
# params: (out, in) layout — I'm choosing the
# same shape convention as nn.Linear.weight here, not nn.Linear(in, out) 
# i.e. more maths-like
def init_params(key):
	# standard pattern in JAX to keep each random client independent
	k1, k2 = jax.random.split(key) 
	
	return {
		"W1": jax.random.normal(k1, (32, 10)) * 0.1,
		"b1": jnp.zeros(32),
		"W2": jax.random.normal(k2, (1, 32)) * 0.1,
		"b2": jnp.zeros(1),
	}
	
# notice how we pass params around since they're not stored in a model
# also notice I'm using maths conventions for dimensions here
def forward_single(params, x):
	h = nn.relu(params["W1"] @ x + params["b1"])
	return params["W2"] @ h + params["b2"]
	
# `vmap` (see JAX docs) enables us to vectorise a function so it's fast on batches
# `in_axes` tells `vmap` which axis of each argument is the "batch" axis to map over — and which arguments have no batch axis at all. lines up with forward_single's inputs
# i.e. expects x to be (batch_size, dimensions) and will vectorise over the batch
forward = jax.vmap(forward_single, in_axes=(None, 0))

# need our own loss as opposed to PyTorch `nn.MSELoss()`
def mse_loss(params, batch_x, batch_y):
	preds = forward(params, batch_x)
	return jnp.mean((preds - batch_y) ** 2)
	
#
# Actually train
#

params = init_params(jax.random.key(42)) # 42 is our seed
optimiser = optax.adam(learning_rate=0.01)
# state as always is immutable, so we update it by creating new values
# functions are just logic, so to speak
opt_state = optimiser.init(params)

# JIT compile functions that are reused a lot, for performance
@jax.jit
def train_step(params, opt_state, batch_x, batch_y):
	# `jax.value_and_grad` is built-in and so valuable!
	loss, grads = jax.value_and_grad(mse_loss)(params, batch_x, batch_y)
	
	# I guess see `optax` docs for this flow
	updates, opt_state = optimiser.update(grads, opt_state, params)
	params = optax.apply_updates(params, updates)
	
	# because we can't mutate state, we return new state
	# return loss in case we wanted to monitor it
	return params, opt_state, loss
	
	
	
# ! heads up: batch_x, batch_y are torch.Tensor here (from the PyTorch DataLoader),
# not jnp.ndarray — this works because JAX implicitly converts anything
# exposing the array interface (works for CPU tensors; GPU tensors or
# tensors with requires_grad=True won't convert cleanly).
# safer/more explicit: batch_x, batch_y = jnp.asarray(batch_x), jnp.asarray(batch_y)
	
# regular training loop
for epoch in range(5):
	for batch_x, batch_y in dataloader:
		# better: batch_x, batch_y = jnp.asarray(batch_x), jnp.asarray(batch_y)
		params, opt_state, _ = train_step(params, opt_state, batch_x, batch_y)

# JAX with Flax

# the more recent module `nnx`. `jax.linen` etc may still be around in production
from flax import nnx 
# already imported:
#import jax.numpy as jnp 
#import optax


# the linear layers are now stateful, managed with 
# jax but under the hood
class MLP(nnx.Module):
	def __init__(self, rngs: nnx.Rngs):
		# ! notice we're back to PyTorch dimension notation instead of 
		# maths convention
		# Also note state is managed by `flax` under the hood here.	
		self.linear1 = nnx.Linear(10, 32, rngs=rngs)
		self.linear2 = nnx.Linear(32, 1, rngs=rngs)
		
	def __call__(self, x):
		x = nnx.relu(self.linear1(x))
		return self.linear2(x)
		
		
model = MLP(rngs=nnx.Rngs(42))
# the optimiser wraps the model and holds a reference to it — like torch.optim 
optimiser = nnx.Optimizer(model, optax.adam(1e-2), wrt=nnx.Param)

# ! @jax.jit only knows how to handle raw JAX i.e. plain pytrees
# so we need to use the `nnx` version that can deal with the `Module`.
# 
# Also we have loss_fn inside train_step just for argument scoping convienience.
# And we define a train_step in JAX/Flax examples so we can use a JIT decorator.
@nnx.jit
def train_step(model, optimiser, batch_x, batch_y):
	def loss_fn(model):
		preds = model(batch_x)
		return jnp.mean((preds - batch_y) ** 2)
		
	loss, grads = nnx.value_and_grad(loss_fn)(model)
	optimiser.update(model, grads)
	
	return loss
	
for epoch in range(5):
	for batch_x, batch_y in dataloader:
		# better: batch_x, batch_y = jnp.asarray(batch_x), jnp.asarray(batch_y)
		_ = train_step(model, optimiser, batch_x, batch_y)