class MLP(nn.Module):\n",
" def __init__(self, in_features, hidden_sizes, out_features, *, rngs):\n",
" super().__init__()\n",
" widths = [in_features, *hidden_sizes]\n",
" self.hidden = nn.ModuleList(\n",
" [nn.Linear(a, b, rngs=rngs) for a, b in zip(widths[:-1], widths[1:], strict=True)]\n",
" )\n",
" self.head = nn.Linear(widths[-1], out_features, rngs=rngs)\n",
"\n",
" def forward(self, x):\n",
" for layer in self.hidden:\n",
" x = F.relu(layer(x))\n",
" return self.head(x)\n",
"@jit\n",
"def classifier_step(model, optimizer, batch):\n",
" x, y = batch\n",
"\n",
" def objective(candidate):\n",
" logits = candidate(x)\n",
" return F.cross_entropy(logits, y), (logits.argmax(-1) == y).sum()\n",
"\n",
" (loss, correct), grads = value_and_grad(objective, has_aux=True)(model)\n",
" optimizer.step(grads)\n",
" return {\"loss\": loss, \"correct\": correct, \"count\": y.size}\n",
"