exercise2_solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z5lAxbrS145P"
   },
   "source": [
    "### Exercise 2 Solution\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/rastringer/jax_notebooks/blob/master/exercise2_solution.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 3269,
     "status": "ok",
     "timestamp": 1685486924860,
     "user": {
      "displayName": "Robin Stringer",
      "userId": "05796722230835218202"
     },
     "user_tz": -60
    },
    "id": "jaAZzJrN1_2p",
    "outputId": "f36117e5-a194-49d6-cad4-bac232275010"
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import numpy as jnp, random, lax, jit\n",
    "from flax import linen as nn\n",
    "\n",
    "X = jnp.ones((1, 10))\n",
    "Y = jnp.ones((5,))\n",
    "\n",
    "model = nn.Dense(features=5)\n",
    "\n",
    "@jit\n",
    "def predict(params):\n",
    "  return model.apply({'params': params}, X)\n",
    "\n",
    "@jit\n",
    "def loss_fn(params):\n",
    "  return jnp.mean(jnp.abs(Y - predict(params)))\n",
    "\n",
    "@jit\n",
    "def init_params(rng):\n",
    "  mlp_variables = model.init({'params': rng}, X)\n",
    "  return mlp_variables['params']\n",
    "\n",
    "# Get initial parameters\n",
    "params = init_params(jax.random.PRNGKey(42))\n",
    "print(\"initial params\", params)\n",
    "\n",
    "# Run SGD.\n",
    "for i in range(50):\n",
    "  loss, grad = jax.value_and_grad(loss_fn)(params)\n",
    "  print(i, \"loss = \", loss, \"Yhat = \", predict(params))\n",
    "  lr = 0.03\n",
    "  params = jax.tree_util.tree_map(lambda x, d: x - lr * d, params, grad)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyP2iG/ZUvKsGUS/phPkkvLs",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python (Local)",
   "language": "python",
   "name": "local-base"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}