docs/notebooks/optax_update_guide.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dHMnJTK9R5n9"
   },
   "source": [
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/notebooks/optax_update_guide.ipynb)\n",
    "[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/notebooks/optax_update_guide.ipynb)\n",
    "\n",
    "Colab for\n",
    "https://flax.readthedocs.io/en/latest/howtos/optax_update_guide.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fCCY-S009eHv"
   },
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "I4PiwrnnO6Fw",
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "!pip install -q flax optax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "7hDWlLOOt4U6"
   },
   "outputs": [],
   "source": [
    "from typing import Sequence\n",
    "\n",
    "import flax\n",
    "from  flax.training import train_state\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import flax.linen as nn\n",
    "import optax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mb2xGRAwueSa",
    "outputId": "30260605-6773-482d-be0e-0383e63b9fa2"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "batch = {\n",
    "    'image': jnp.ones([1, 28, 28, 1]),\n",
    "    'label': jnp.array([0]),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lf11Nzj-t32w",
    "outputId": "c54a570b-d76a-43bb-f1ab-5cbbbfd6f584"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FrozenDict({\n",
       "    params: {\n",
       "        Dense_0: {\n",
       "            bias: (50,),\n",
       "            kernel: (784, 50),\n",
       "        },\n",
       "        Dense_1: {\n",
       "            bias: (10,),\n",
       "            kernel: (50, 10),\n",
       "        },\n",
       "    },\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Perceptron(nn.Module):\n",
    "  units: Sequence[int]\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "    x = x.reshape([x.shape[0], -1]) / 255.\n",
    "    x = nn.Dense(50)(x)\n",
    "    x = nn.relu(x)\n",
    "    return nn.Dense(10)(x)\n",
    "\n",
    "def loss(params, batch):\n",
    "  logits = model.apply({'params': params}, batch['image'])\n",
    "  one_hot = jax.nn.one_hot(batch['label'], 10)\n",
    "  return jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))\n",
    "\n",
    "model = Perceptron([50, 10])\n",
    "variables = model.init(jax.random.PRNGKey(0), batch['image'])\n",
    "\n",
    "jax.tree_map(jnp.shape, variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xKm0nn4X57Vg",
    "outputId": "e1794ea2-ea91-45b5-8b89-d39d2d923cc5"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tensorflow_datasets/core/dataset_builder.py:598: get_single_element (from tensorflow.python.data.experimental.ops.get_single_element) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.Dataset.get_single_element()`.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'image': (128, 28, 28, 1), 'label': (128,)}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow_datasets as tfds\n",
    "\n",
    "builder = tfds.builder('mnist')\n",
    "builder.download_and_prepare()\n",
    "ds_test = jax.tree_map(jnp.array, builder.as_dataset('test', batch_size=-1))\n",
    "get_ds_train = lambda: (\n",
    "    jax.tree_map(jnp.array, x)\n",
    "    for x in builder.as_dataset('train').batch(128))\n",
    "batch = next(get_ds_train())\n",
    "jax.tree_map(jnp.shape, batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eCceGZ_Kvko5",
    "outputId": "af1f8657-b3eb-4c9b-d073-48954481166b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.103, dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def eval(params):\n",
    "  logits = model.apply({'params': params}, ds_test['image'])\n",
    "  return (logits.argmax(axis=-1) == ds_test['label']).mean()\n",
    "\n",
    "eval(variables['params'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "rqQiq3ugxKjX"
   },
   "outputs": [],
   "source": [
    "learning_rate, momentum = 0.01, 0.9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JyXnY2WW9gfR"
   },
   "source": [
    "### Replacing `flax.optim` with `optax`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7mlz-P5AwBc2",
    "outputId": "05cb4271-e407-4798-d0ea-cd8dfbd9f2a1"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.9165, dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(optimizer, batch):\n",
    "  grads = jax.grad(loss)(optimizer.target, batch)\n",
    "  return optimizer.apply_gradient(grads)\n",
    "\n",
    "optimizer = flax.optim.Momentum(learning_rate, momentum).create(\n",
    "    variables['params'])\n",
    "for batch in get_ds_train():\n",
    "  optimizer = train_step(optimizer, batch)\n",
    "\n",
    "eval(optimizer.target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "n9IxglwJxD3X",
    "outputId": "a42f9e8d-d1fe-4221-ec28-b7258174b16c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.9165, dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tx = optax.sgd(learning_rate, momentum)\n",
    "params = variables['params']\n",
    "opt_state = tx.init(params)\n",
    "\n",
    "@jax.jit\n",
    "def train_step(params, opt_state, batch):\n",
    "  grads = jax.grad(loss)(params, batch)\n",
    "  updates, opt_state = tx.update(grads, opt_state)\n",
    "  params = optax.apply_updates(params, updates)\n",
    "  return params, opt_state\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  params, opt_state = train_step(params, opt_state, batch)\n",
    "\n",
    "eval(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mBcp17BEvBVs",
    "outputId": "dd912efc-669f-42c2-86b2-e242cecc67ce"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.9165, dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(state, batch):\n",
    "  def loss(params):\n",
    "    logits = state.apply_fn({'params': params}, batch['image'])\n",
    "    one_hot = jax.nn.one_hot(batch['label'], 10)\n",
    "    return jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))\n",
    "  grads = jax.grad(loss)(state.params)\n",
    "  return state.apply_gradients(grads=grads)\n",
    "\n",
    "tx = optax.sgd(learning_rate, momentum)\n",
    "state = train_state.TrainState.create(\n",
    "    apply_fn=model.apply, tx=tx, params=variables['params'],\n",
    ")\n",
    "opt_state = tx.init(params)\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  state = train_step(state, batch)\n",
    "\n",
    "eval(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4ute1zBpRnaq"
   },
   "source": [
    "### Composable Gradient Transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "M2WjJ7HT8GMn",
    "outputId": "7ae7117e-5d1d-44ca-abb8-65b0db2c0eb6"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.9165, dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(params, opt_state, batch):\n",
    "  grads = jax.grad(loss)(params, batch)\n",
    "  updates, opt_state = tx.update(grads, opt_state)\n",
    "  params = optax.apply_updates(params, updates)\n",
    "  return params, opt_state\n",
    "\n",
    "tx = optax.chain(\n",
    "    optax.trace(decay=momentum),\n",
    "    optax.scale(-learning_rate),\n",
    ")\n",
    "params = variables['params']\n",
    "opt_state = tx.init(params)\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  params, opt_state = train_step(params, opt_state, batch)\n",
    "\n",
    "eval(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vFt96TU-rSQM"
   },
   "source": [
    "### Weight Decay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "Qbq9vK24-omQ"
   },
   "outputs": [],
   "source": [
    "weight_decay = 1e-5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cx1YCFVL9ktA",
    "outputId": "bb6795a7-c5c2-458a-d3e9-4b769b857fed"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.95129997, dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(optimizer, batch):\n",
    "  grads = jax.grad(loss)(optimizer.target, batch)\n",
    "  return optimizer.apply_gradient(grads)\n",
    "\n",
    "optimizer = flax.optim.Adam(learning_rate, weight_decay=weight_decay).create(\n",
    "    variables['params'])\n",
    "for batch in get_ds_train():\n",
    "  optimizer = train_step(optimizer, batch)\n",
    "\n",
    "eval(optimizer.target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kIWCQz33-p98",
    "outputId": "84753c79-314b-4982-97e6-0c61a490eab1"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.9517, dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(params, opt_state, batch):\n",
    "  grads = jax.grad(loss)(params, batch)\n",
    "  updates, opt_state = tx.update(grads, opt_state, params)\n",
    "  params = optax.apply_updates(params, updates)\n",
    "  return params, opt_state\n",
    "\n",
    "tx = optax.chain(\n",
    "    optax.scale_by_adam(),\n",
    "    optax.add_decayed_weights(weight_decay),\n",
    "    optax.scale(-learning_rate),\n",
    ")\n",
    "params = variables['params']\n",
    "opt_state = tx.init(params)\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  params, opt_state = train_step(params, opt_state, batch)\n",
    "\n",
    "eval(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MZE97FEbCgmn"
   },
   "source": [
    "### Gradient Clipping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "Y_DCoogODZL6"
   },
   "outputs": [],
   "source": [
    "grad_clip_norm = 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mFnX8fb3Chwb",
    "outputId": "aae283d2-623e-4a43-c001-3e62b11483ae"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.91679996, dtype=float32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(optimizer, batch):\n",
    "  grads = jax.grad(loss)(optimizer.target, batch)\n",
    "  grads_flat, _ = jax.tree_flatten(grads)\n",
    "  global_l2 = jnp.sqrt(sum([jnp.vdot(p, p) for p in grads_flat]))\n",
    "  g_factor = jnp.minimum(1.0, grad_clip_norm / global_l2)\n",
    "  grads = jax.tree_map(lambda g: g * g_factor, grads)\n",
    "  return optimizer.apply_gradient(grads)\n",
    "\n",
    "optimizer = flax.optim.Momentum(learning_rate, momentum).create(\n",
    "    variables['params'])\n",
    "for batch in get_ds_train():\n",
    "  optimizer = train_step(optimizer, batch)\n",
    "\n",
    "eval(optimizer.target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aJYN2A-TDhp3",
    "outputId": "a9563d3a-7dc6-4ecf-bb7b-49356cf9fe13"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.91679996, dtype=float32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(params, opt_state, batch):\n",
    "  grads = jax.grad(loss)(params, batch)\n",
    "  updates, opt_state = tx.update(grads, opt_state)\n",
    "  params = optax.apply_updates(params, updates)\n",
    "  return params, opt_state\n",
    "\n",
    "tx = optax.chain(\n",
    "    optax.clip_by_global_norm(grad_clip_norm),\n",
    "    optax.trace(decay=momentum),\n",
    "    optax.scale(-learning_rate),\n",
    ")\n",
    "params = variables['params']\n",
    "opt_state = tx.init(params)\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  params, opt_state = train_step(params, opt_state, batch)\n",
    "\n",
    "eval(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d9e9YmNHE-xV"
   },
   "source": [
    "### Learning Rate Schedules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "zCqz6fDsFB5n"
   },
   "outputs": [],
   "source": [
    "schedule = lambda step: learning_rate * jnp.exp(step * 1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "NinuzivVFYb5",
    "outputId": "c90b880e-d6f0-40fc-94b8-28d0622d8440"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.9201, dtype=float32)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(step, optimizer, batch):\n",
    "  grads = jax.grad(loss)(optimizer.target, batch)\n",
    "  return step + 1, optimizer.apply_gradient(grads, learning_rate=schedule(step))\n",
    "\n",
    "optimizer = flax.optim.Momentum(learning_rate, momentum).create(\n",
    "    variables['params'])\n",
    "step = jnp.array(0)\n",
    "for batch in get_ds_train():\n",
    "  step, optimizer = train_step(step, optimizer, batch)\n",
    "\n",
    "eval(optimizer.target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lzaYwXzuFp-L",
    "outputId": "0ed7f41f-cc2d-46c6-ae9f-4be5b906fb7f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.9201, dtype=float32)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(params, opt_state, batch):\n",
    "  grads = jax.grad(loss)(params, batch)\n",
    "  updates, opt_state = tx.update(grads, opt_state)\n",
    "  params = optax.apply_updates(params, updates)\n",
    "  return params, opt_state\n",
    "\n",
    "tx = optax.chain(\n",
    "    optax.trace(decay=momentum),\n",
    "    optax.scale_by_schedule(lambda step: -schedule(step)),\n",
    ")\n",
    "params = variables['params']\n",
    "opt_state = tx.init(params)\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  params, opt_state = train_step(params, opt_state, batch)\n",
    "\n",
    "eval(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HLkQDKK0GCHH"
   },
   "source": [
    "### Multiple Optimizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "d-2veL8lGDbV",
    "outputId": "b006d3cc-eff6-410a-e201-ccd9464db9d7"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.91679996, dtype=float32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(optimizer, batch):\n",
    "  grads = jax.grad(loss)(optimizer.target, batch)\n",
    "  return optimizer.apply_gradient(grads)\n",
    "\n",
    "kernels = flax.traverse_util.ModelParamTraversal(lambda p, _: 'kernel' in p)\n",
    "biases = flax.traverse_util.ModelParamTraversal(lambda p, _: 'bias' in p)\n",
    "kernel_opt = flax.optim.Momentum(learning_rate, momentum)\n",
    "bias_opt = flax.optim.Momentum(learning_rate * 0.1, momentum)\n",
    "optimizer = flax.optim.MultiOptimizer(\n",
    "    (kernels, kernel_opt),\n",
    "    (biases, bias_opt)\n",
    ").create(variables['params'])\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  optimizer = train_step(optimizer, batch)\n",
    "\n",
    "eval(optimizer.target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MvQlkiCuHr41",
    "outputId": "88b4b0dd-4b80-4c5b-c21d-803c097b8cc6"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.91679996, dtype=float32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(params, opt_state, batch):\n",
    "  grads = jax.grad(loss)(params, batch)\n",
    "  updates, opt_state = tx.update(grads, opt_state)\n",
    "  params = optax.apply_updates(params, updates)\n",
    "  return params, opt_state\n",
    "\n",
    "kernels = flax.traverse_util.ModelParamTraversal(lambda p, _: 'kernel' in p)\n",
    "biases = flax.traverse_util.ModelParamTraversal(lambda p, _: 'bias' in p)\n",
    "\n",
    "all_false = jax.tree_map(lambda _: False, params)\n",
    "kernels_mask = kernels.update(lambda _: True, all_false)\n",
    "biases_mask = biases.update(lambda _: True, all_false)\n",
    "\n",
    "tx = optax.chain(\n",
    "    optax.trace(decay=momentum),\n",
    "    optax.masked(optax.scale(-learning_rate), kernels_mask),\n",
    "    optax.masked(optax.scale(-learning_rate * 0.1), biases_mask),\n",
    ")\n",
    "params = variables['params']\n",
    "opt_state = tx.init(params)\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  params, opt_state = train_step(params, opt_state, batch)\n",
    "\n",
    "eval(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "v2omX108JYDo",
    "outputId": "a9f45c8c-0db5-4b5e-b429-e6d79b22eca0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.91679996, dtype=float32)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def train_step(params, opt_state, batch):\n",
    "  grads = jax.grad(loss)(params, batch)\n",
    "  updates, opt_state = tx.update(grads, opt_state)\n",
    "  params = optax.apply_updates(params, updates)\n",
    "  return params, opt_state\n",
    "\n",
    "kernels = flax.traverse_util.ModelParamTraversal(lambda p, _: 'kernel' in p)\n",
    "biases = flax.traverse_util.ModelParamTraversal(lambda p, _: 'bias' in p)\n",
    "\n",
    "all_false = jax.tree_map(lambda _: False, params)\n",
    "kernels_mask = kernels.update(lambda _: True, all_false)\n",
    "biases_mask = biases.update(lambda _: True, all_false)\n",
    "\n",
    "tx = optax.chain(\n",
    "    optax.trace(decay=momentum),\n",
    "    optax.multi_transform({\n",
    "      'kernels': optax.scale(-learning_rate),\n",
    "      'biases': optax.scale(-learning_rate * 0.1),\n",
    "  }, kernels.update(lambda _: 'kernels',\n",
    "                    biases.update(lambda _: 'biases', params))),\n",
    ")\n",
    "params = variables['params']\n",
    "opt_state = tx.init(params)\n",
    "\n",
    "for batch in get_ds_train():\n",
    "  params, opt_state = train_step(params, opt_state, batch)\n",
    "\n",
    "eval(params)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Optax Update Guide",
   "provenance": [],
   "toc_visible": true
  },
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}