{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_"
      },
      "source": [
        "# Visualization\n",
        "\n",
        "Haiku supports two ways to visualize your program. To use these you need to install two additional dependencies:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "_"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: dm-tree in /tmp/haiku-env/lib/python3.11/site-packages (0.1.8)\r\n",
            "Requirement already satisfied: graphviz in /tmp/haiku-env/lib/python3.11/site-packages (0.20.1)\r\n"
          ]
        }
      ],
      "source": [
        "!pip install dm-tree graphviz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "_"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import haiku as hk"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_"
      },
      "source": [
        "## Tabulate\n",
        "\n",
        "Like many neural network libraries, Haiku supports showing a summary of the execution of your program as a table of modules. Haiku's approach is to trace the execution of your program and to produce a table of (interesting) module method calls.\n",
        "\n",
        "For example, the interesting methods for a 3 layer MLP would be `MLP.__call__` which in turns calls `Linear.__call__` on three inner modules. For each module method we show columns relating to the input/output size of arrays, as well as details of the modules parameters and where it fits in the module hierarchy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "_"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "+-------------------------+------------------------------------------+-----------------+------------+------------+---------------+---------------+\n",
            "| Module                  | Config                                   | Module params   | Input      | Output     |   Param count |   Param bytes |\n",
            "+=========================+==========================================+=================+============+============+===============+===============+\n",
            "| mlp (MLP)               | MLP(output_sizes=[300, 100, 10])         |                 | f32[8,784] | f32[8,10]  |       266,610 |       1.07 MB |\n",
            "+-------------------------+------------------------------------------+-----------------+------------+------------+---------------+---------------+\n",
            "| mlp/~/linear_0 (Linear) | Linear(output_size=300, name='linear_0') | w: f32[784,300] | f32[8,784] | f32[8,300] |       235,500 |     942.00 KB |\n",
            "|  └ mlp (MLP)            |                                          | b: f32[300]     |            |            |               |               |\n",
            "+-------------------------+------------------------------------------+-----------------+------------+------------+---------------+---------------+\n",
            "| mlp/~/linear_1 (Linear) | Linear(output_size=100, name='linear_1') | w: f32[300,100] | f32[8,300] | f32[8,100] |        30,100 |     120.40 KB |\n",
            "|  └ mlp (MLP)            |                                          | b: f32[100]     |            |            |               |               |\n",
            "+-------------------------+------------------------------------------+-----------------+------------+------------+---------------+---------------+\n",
            "| mlp/~/linear_2 (Linear) | Linear(output_size=10, name='linear_2')  | w: f32[100,10]  | f32[8,100] | f32[8,10]  |         1,010 |       4.04 KB |\n",
            "|  └ mlp (MLP)            |                                          | b: f32[10]      |            |            |               |               |\n",
            "+-------------------------+------------------------------------------+-----------------+------------+------------+---------------+---------------+\n"
          ]
        }
      ],
      "source": [
        "def f(x):\n",
        "  return hk.nets.MLP([300, 100, 10])(x)\n",
        "\n",
        "f = hk.transform(f)\n",
        "x = jnp.ones([8, 28 * 28])\n",
        "\n",
        "print(hk.experimental.tabulate(f)(x))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_"
      },
      "source": [
        "We also offer access to the raw data used to build this table if you want to create your own summary:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "_"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "MethodInvocation(module_details=ModuleDetails(module=MLP(output_sizes=[300, 100, 10]), method_name='__call__', params={'mlp/~/linear_0/b': f32[300], 'mlp/~/linear_0/w': f32[784,300], 'mlp/~/linear_1/b': f32[100], 'mlp/~/linear_1/w': f32[300,100], 'mlp/~/linear_2/b': f32[10], 'mlp/~/linear_2/w': f32[100,10]}, state={}), args_spec=(f32[8,784],), kwargs_spec={}, output_spec=f32[8,10], context=MethodContext(module=MLP(output_sizes=[300, 100, 10]), method_name='__call__', orig_method=functools.partial(<function MLP.__call__ at 0x7f173d83f600>, MLP(output_sizes=[300, 100, 10])), orig_class=<class 'haiku._src.nets.mlp.MLP'>), call_stack=(ModuleDetails(module=MLP(output_sizes=[300, 100, 10]), method_name='__call__', params={'mlp/~/linear_0/b': f32[300], 'mlp/~/linear_0/w': f32[784,300], 'mlp/~/linear_1/b': f32[100], 'mlp/~/linear_1/w': f32[300,100], 'mlp/~/linear_2/b': f32[10], 'mlp/~/linear_2/w': f32[100,10]}, state={}),))\n",
            "MethodInvocation(module_details=ModuleDetails(module=Linear(output_size=300, name='linear_0'), method_name='__call__', params={'mlp/~/linear_0/b': f32[300], 'mlp/~/linear_0/w': f32[784,300]}, state={}), args_spec=(f32[8,784],), kwargs_spec={}, output_spec=f32[8,300], context=MethodContext(module=Linear(output_size=300, name='linear_0'), method_name='__call__', orig_method=functools.partial(<function Linear.__call__ at 0x7f173d927e20>, Linear(output_size=300, name='linear_0')), orig_class=<class 'haiku._src.basic.Linear'>), call_stack=(ModuleDetails(module=Linear(output_size=300, name='linear_0'), method_name='__call__', params={'mlp/~/linear_0/b': f32[300], 'mlp/~/linear_0/w': f32[784,300]}, state={}), ModuleDetails(module=MLP(output_sizes=[300, 100, 10]), method_name='__call__', params={'mlp/~/linear_0/b': f32[300], 'mlp/~/linear_0/w': f32[784,300], 'mlp/~/linear_1/b': f32[100], 'mlp/~/linear_1/w': f32[300,100], 'mlp/~/linear_2/b': f32[10], 'mlp/~/linear_2/w': f32[100,10]}, state={})))\n",
            "MethodInvocation(module_details=ModuleDetails(module=Linear(output_size=100, name='linear_1'), method_name='__call__', params={'mlp/~/linear_1/b': f32[100], 'mlp/~/linear_1/w': f32[300,100]}, state={}), args_spec=(f32[8,300],), kwargs_spec={}, output_spec=f32[8,100], context=MethodContext(module=Linear(output_size=100, name='linear_1'), method_name='__call__', orig_method=functools.partial(<function Linear.__call__ at 0x7f173d927e20>, Linear(output_size=100, name='linear_1')), orig_class=<class 'haiku._src.basic.Linear'>), call_stack=(ModuleDetails(module=Linear(output_size=100, name='linear_1'), method_name='__call__', params={'mlp/~/linear_1/b': f32[100], 'mlp/~/linear_1/w': f32[300,100]}, state={}), ModuleDetails(module=MLP(output_sizes=[300, 100, 10]), method_name='__call__', params={'mlp/~/linear_0/b': f32[300], 'mlp/~/linear_0/w': f32[784,300], 'mlp/~/linear_1/b': f32[100], 'mlp/~/linear_1/w': f32[300,100], 'mlp/~/linear_2/b': f32[10], 'mlp/~/linear_2/w': f32[100,10]}, state={})))\n",
            "MethodInvocation(module_details=ModuleDetails(module=Linear(output_size=10, name='linear_2'), method_name='__call__', params={'mlp/~/linear_2/b': f32[10], 'mlp/~/linear_2/w': f32[100,10]}, state={}), args_spec=(f32[8,100],), kwargs_spec={}, output_spec=f32[8,10], context=MethodContext(module=Linear(output_size=10, name='linear_2'), method_name='__call__', orig_method=functools.partial(<function Linear.__call__ at 0x7f173d927e20>, Linear(output_size=10, name='linear_2')), orig_class=<class 'haiku._src.basic.Linear'>), call_stack=(ModuleDetails(module=Linear(output_size=10, name='linear_2'), method_name='__call__', params={'mlp/~/linear_2/b': f32[10], 'mlp/~/linear_2/w': f32[100,10]}, state={}), ModuleDetails(module=MLP(output_sizes=[300, 100, 10]), method_name='__call__', params={'mlp/~/linear_0/b': f32[300], 'mlp/~/linear_0/w': f32[784,300], 'mlp/~/linear_1/b': f32[100], 'mlp/~/linear_1/w': f32[300,100], 'mlp/~/linear_2/b': f32[10], 'mlp/~/linear_2/w': f32[100,10]}, state={})))\n"
          ]
        }
      ],
      "source": [
        "for method_invocation in hk.experimental.eval_summary(f)(x):\n",
        "  print(method_invocation)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_"
      },
      "source": [
        "## Graphviz (aka. `to_dot`)\n",
        "\n",
        "Haiku supports rendering your program as a graphviz graph. We show all of the JAX primitives involved in a given computation clustered by Haiku module.\n",
        "\n",
        "Lets start by visualizing a simple program not using Haiku modules:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "_"
      },
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<?xml-stylesheet href=\"\n",
              "  data:text/css,\n",
              "  @import url(https://fonts.googleapis.com/css?family=Roboto:400,700);\n",
              "  svg text {\n",
              "    font-family: 'Roboto';\n",
              "  }\n",
              "  .node text {\n",
              "    font-size: 12px;\n",
              "  }\n",
              "%23node486529600:hover { stroke-width: 0.2em; }\n",
              "%23node486529600:hover ~ %23node486529601 { stroke-width: 0.2em; }\n",
              "%23node486529601:hover ~ %23node486529600 { stroke-width: 0.2em; }\n",
              "%23node486529601:hover { stroke-width: 0.2em; }\n",
              "\" type=\"text/css\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 2.43.0 (0)\n",
              " -->\n",
              "<!-- Title: G Pages: 1 -->\n",
              "<svg width=\"138pt\" height=\"389pt\"\n",
              " viewBox=\"0.00 0.00 138.00 389.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 385)\">\n",
              "<title>G</title>\n",
              "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-385 134,-385 134,4 -4,4\"/>\n",
              "<text text-anchor=\"start\" x=\"61.5\" y=\"-362.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"19.00\">f</text>\n",
              "<!-- 49830512 -->\n",
              "<g id=\"node49830512\" class=\"node\">\n",
              "<title>49830512</title>\n",
              "<g id=\"a_node49830512\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"128.5,-280 74.5,-280 74.5,-244 128.5,-244 128.5,-280\"/>\n",
              "<text text-anchor=\"start\" x=\"93\" y=\"-265.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">sin</text>\n",
              "<text text-anchor=\"start\" x=\"84.5\" y=\"-251.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[1]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 51865824 -->\n",
              "<g id=\"node51865824\" class=\"node\">\n",
              "<title>51865824</title>\n",
              "<g id=\"a_node51865824\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"101.5,-208 47.5,-208 47.5,-172 101.5,-172 101.5,-208\"/>\n",
              "<text text-anchor=\"start\" x=\"65\" y=\"-193.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">cos</text>\n",
              "<text text-anchor=\"start\" x=\"57.5\" y=\"-179.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[1]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 49830512&#45;&gt;51865824 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>49830512&#45;&gt;51865824</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M94.83,-243.7C91.79,-235.81 88.12,-226.3 84.74,-217.55\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"87.96,-216.17 81.1,-208.1 81.43,-218.69 87.96,-216.17\"/>\n",
              "</g>\n",
              "<!-- 52344992 -->\n",
              "<g id=\"node52344992\" class=\"node\">\n",
              "<title>52344992</title>\n",
              "<g id=\"a_node52344992\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"129.5,-129 75.5,-129 75.5,-93 129.5,-93 129.5,-129\"/>\n",
              "<text text-anchor=\"start\" x=\"91.5\" y=\"-114.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">add</text>\n",
              "<text text-anchor=\"start\" x=\"85.5\" y=\"-100.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[1]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 49830512&#45;&gt;52344992 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>49830512&#45;&gt;52344992</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M105.15,-243.91C107.18,-233.57 109.48,-220.09 110.5,-208 112.45,-184.91 109.94,-158.68 107.23,-139.49\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"110.66,-138.7 105.69,-129.34 103.73,-139.76 110.66,-138.7\"/>\n",
              "</g>\n",
              "<!-- 51865824&#45;&gt;52344992 -->\n",
              "<g id=\"edge4\" class=\"edge\">\n",
              "<title>51865824&#45;&gt;52344992</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M80.71,-171.91C84.27,-162.14 88.8,-149.68 92.81,-138.65\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"96.18,-139.64 96.3,-129.04 89.6,-137.24 96.18,-139.64\"/>\n",
              "</g>\n",
              "<!-- 49103760 -->\n",
              "<g id=\"node49103760\" class=\"node\">\n",
              "<title>49103760</title>\n",
              "<g id=\"a_node49103760\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#aabfff\" stroke=\"#0053d6\" stroke-width=\"2\" points=\"93,-50 38,-50 38,0 93,0 93,-50\"/>\n",
              "<text text-anchor=\"start\" x=\"46\" y=\"-35.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">output</text>\n",
              "<text text-anchor=\"start\" x=\"54.5\" y=\"-21.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">add</text>\n",
              "<text text-anchor=\"start\" x=\"48.5\" y=\"-7.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[1]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 52344992&#45;&gt;49103760 -->\n",
              "<g id=\"edge6\" class=\"edge\">\n",
              "<title>52344992&#45;&gt;49103760</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M94.83,-92.6C90.6,-82.98 85.2,-70.72 80.19,-59.34\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"83.34,-57.81 76.1,-50.07 76.93,-60.63 83.34,-57.81\"/>\n",
              "</g>\n",
              "<!-- 486529600 -->\n",
              "<g id=\"node486529600\" class=\"node\">\n",
              "<title>486529600</title>\n",
              "<polygon fill=\"#ffdeaf\" stroke=\"#ff8a4f\" points=\"130,-352 73,-352 73,-316 130,-316 130,-352\"/>\n",
              "<text text-anchor=\"start\" x=\"81\" y=\"-337.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">args[0]</text>\n",
              "<text text-anchor=\"start\" x=\"84.5\" y=\"-323.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[1]</text>\n",
              "</g>\n",
              "<!-- 486529600&#45;&gt;49830512 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>486529600&#45;&gt;49830512</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M101.5,-315.7C101.5,-307.98 101.5,-298.71 101.5,-290.11\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"105,-290.1 101.5,-280.1 98,-290.1 105,-290.1\"/>\n",
              "</g>\n",
              "<!-- 486529601 -->\n",
              "<g id=\"node486529601\" class=\"node\">\n",
              "<title>486529601</title>\n",
              "<polygon fill=\"#ffeacc\" stroke=\"#ff8a4f\" points=\"57,-136 0,-136 0,-86 57,-86 57,-136\"/>\n",
              "<text text-anchor=\"start\" x=\"8.5\" y=\"-121.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\" fill=\"#565858\">(reuse)</text>\n",
              "<text text-anchor=\"start\" x=\"8\" y=\"-107.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\" fill=\"#565858\">args[0]</text>\n",
              "<text text-anchor=\"start\" x=\"11.5\" y=\"-93.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#565858\">f32[1]</text>\n",
              "</g>\n",
              "<!-- 486529601&#45;&gt;49103760 -->\n",
              "<g id=\"edge5\" class=\"edge\">\n",
              "<title>486529601&#45;&gt;49103760</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M39.2,-85.7C42.81,-77.52 46.9,-68.24 50.76,-59.47\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"53.99,-60.8 54.82,-50.24 47.59,-57.98 53.99,-60.8\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x7f1800a33bd0>"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "def f(a):\n",
        "  b = jnp.sin(a)\n",
        "  c = jnp.cos(b)\n",
        "  d = b + c\n",
        "  e = a + d\n",
        "  return e\n",
        "\n",
        "x = jnp.ones([1])\n",
        "dot = hk.to_dot(f)(x)\n",
        "\n",
        "import graphviz\n",
        "graphviz.Source(dot)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_"
      },
      "source": [
        "The visualization above shows our program as a simple dataflow graph of our single input highligted in orange (`args[0]`) being passed through some operations and producing a result (highlighted in blue). Primitive operations (e.g. `sin`, `cos` and `add`) are highlighted in yellow.\n",
        "\n",
        "Actual Haiku programs are often far more complex, involving many modules and many more primitive operations. For these programs it is often useful to visualize the program on a module by module basis.\n",
        "\n",
        "`to_dot` offers this by clustering operations by their module. Again it is probably simplest to see an example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "_"
      },
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<?xml-stylesheet href=\"\n",
              "  data:text/css,\n",
              "  @import url(https://fonts.googleapis.com/css?family=Roboto:400,700);\n",
              "  svg text {\n",
              "    font-family: 'Roboto';\n",
              "  }\n",
              "  .node text {\n",
              "    font-size: 12px;\n",
              "  }\n",
              "\" type=\"text/css\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 2.43.0 (0)\n",
              " -->\n",
              "<!-- Title: G Pages: 1 -->\n",
              "<svg width=\"485pt\" height=\"839pt\"\n",
              " viewBox=\"0.00 0.00 485.00 839.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 835)\">\n",
              "<title>G</title>\n",
              "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-835 481,-835 481,4 -4,4\"/>\n",
              "<text text-anchor=\"start\" x=\"167\" y=\"-797.6\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"38.00\">apply_fn</text>\n",
              "<g id=\"clust5\" class=\"cluster\">\n",
              "<title>cluster_139737768113904</title>\n",
              "<path fill=\"#f0f5f5\" stroke=\"#14234b\" d=\"M20,-8C20,-8 457,-8 457,-8 463,-8 469,-14 469,-20 469,-20 469,-761 469,-761 469,-767 463,-773 457,-773 457,-773 20,-773 20,-773 14,-773 8,-767 8,-761 8,-761 8,-20 8,-20 8,-14 14,-8 20,-8\"/>\n",
              "<text text-anchor=\"start\" x=\"215.5\" y=\"-748.4\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"27.00\">mlp</text>\n",
              "</g>\n",
              "<g id=\"clust6\" class=\"cluster\">\n",
              "<title>cluster_139737768115264</title>\n",
              "<path fill=\"#f0f5f5\" stroke=\"#14234b\" d=\"M231,-530C231,-530 449,-530 449,-530 455,-530 461,-536 461,-542 461,-542 461,-671 461,-671 461,-677 455,-683 449,-683 449,-683 231,-683 231,-683 225,-683 219,-677 219,-671 219,-671 219,-542 219,-542 219,-536 225,-530 231,-530\"/>\n",
              "<text text-anchor=\"start\" x=\"280.5\" y=\"-664.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"19.00\">mlp/~/linear_0</text>\n",
              "</g>\n",
              "<g id=\"clust7\" class=\"cluster\">\n",
              "<title>cluster_139737790661312</title>\n",
              "<path fill=\"#f0f5f5\" stroke=\"#14234b\" d=\"M380,-441C380,-441 449,-441 449,-441 455,-441 461,-447 461,-453 461,-453 461,-510 461,-510 461,-516 455,-522 449,-522 449,-522 380,-522 380,-522 374,-522 368,-516 368,-510 368,-510 368,-453 368,-453 368,-447 374,-441 380,-441\"/>\n",
              "<text text-anchor=\"start\" x=\"376\" y=\"-503.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"19.00\">pjit (relu)</text>\n",
              "</g>\n",
              "<g id=\"clust8\" class=\"cluster\">\n",
              "<title>cluster_139737790121520</title>\n",
              "<path fill=\"#f0f5f5\" stroke=\"#14234b\" d=\"M231,-280C231,-280 449,-280 449,-280 455,-280 461,-286 461,-292 461,-292 461,-421 461,-421 461,-427 455,-433 449,-433 449,-433 231,-433 231,-433 225,-433 219,-427 219,-421 219,-421 219,-292 219,-292 219,-286 225,-280 231,-280\"/>\n",
              "<text text-anchor=\"start\" x=\"280.5\" y=\"-414.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"19.00\">mlp/~/linear_1</text>\n",
              "</g>\n",
              "<g id=\"clust9\" class=\"cluster\">\n",
              "<title>cluster_139737767817952</title>\n",
              "<path fill=\"#f0f5f5\" stroke=\"#14234b\" d=\"M380,-191C380,-191 449,-191 449,-191 455,-191 461,-197 461,-203 461,-203 461,-260 461,-260 461,-266 455,-272 449,-272 449,-272 380,-272 380,-272 374,-272 368,-266 368,-260 368,-260 368,-203 368,-203 368,-197 374,-191 380,-191\"/>\n",
              "<text text-anchor=\"start\" x=\"376\" y=\"-253.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"19.00\">pjit (relu)</text>\n",
              "</g>\n",
              "<g id=\"clust10\" class=\"cluster\">\n",
              "<title>cluster_139737768114384</title>\n",
              "<path fill=\"#f0f5f5\" stroke=\"#14234b\" d=\"M99,-16C99,-16 317,-16 317,-16 323,-16 329,-22 329,-28 329,-28 329,-171 329,-171 329,-177 323,-183 317,-183 317,-183 99,-183 99,-183 93,-183 87,-177 87,-171 87,-171 87,-28 87,-28 87,-22 93,-16 99,-16\"/>\n",
              "<text text-anchor=\"start\" x=\"148.5\" y=\"-164.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"19.00\">mlp/~/linear_2</text>\n",
              "</g>\n",
              "<!-- 59282016 -->\n",
              "<g id=\"node59282016\" class=\"node\">\n",
              "<title>59282016</title>\n",
              "<g id=\"a_node59282016\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"452.5,-646 367.5,-646 367.5,-610 452.5,-610 452.5,-646\"/>\n",
              "<text text-anchor=\"start\" x=\"375.5\" y=\"-631.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">dot_general</text>\n",
              "<text text-anchor=\"start\" x=\"381\" y=\"-617.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,300]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 71412688 -->\n",
              "<g id=\"node71412688\" class=\"node\">\n",
              "<title>71412688</title>\n",
              "<g id=\"a_node71412688\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"447,-574 373,-574 373,-538 447,-538 447,-574\"/>\n",
              "<text text-anchor=\"start\" x=\"399\" y=\"-559.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">add</text>\n",
              "<text text-anchor=\"start\" x=\"381\" y=\"-545.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,300]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 59282016&#45;&gt;71412688 -->\n",
              "<g id=\"edge4\" class=\"edge\">\n",
              "<title>59282016&#45;&gt;71412688</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M410,-609.7C410,-601.98 410,-592.71 410,-584.11\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"413.5,-584.1 410,-574.1 406.5,-584.1 413.5,-584.1\"/>\n",
              "</g>\n",
              "<!-- 69211888 -->\n",
              "<g id=\"node69211888\" class=\"node\">\n",
              "<title>69211888</title>\n",
              "<g id=\"a_node69211888\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"349,-646 227,-646 227,-610 349,-610 349,-646\"/>\n",
              "<text text-anchor=\"start\" x=\"235\" y=\"-631.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">broadcast_in_dim</text>\n",
              "<text text-anchor=\"start\" x=\"259\" y=\"-617.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,300]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 69211888&#45;&gt;71412688 -->\n",
              "<g id=\"edge5\" class=\"edge\">\n",
              "<title>69211888&#45;&gt;71412688</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M317.84,-609.88C333.94,-600.64 354,-589.13 371.29,-579.21\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"373.11,-582.2 380.04,-574.19 369.63,-576.13 373.11,-582.2\"/>\n",
              "</g>\n",
              "<!-- 67211248 -->\n",
              "<g id=\"node67211248\" class=\"node\">\n",
              "<title>67211248</title>\n",
              "<g id=\"a_node67211248\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"452,-485 378,-485 378,-449 452,-449 452,-485\"/>\n",
              "<text text-anchor=\"start\" x=\"402\" y=\"-470.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">max</text>\n",
              "<text text-anchor=\"start\" x=\"386\" y=\"-456.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,300]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 71412688&#45;&gt;67211248 -->\n",
              "<g id=\"edge6\" class=\"edge\">\n",
              "<title>71412688&#45;&gt;67211248</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M410.99,-537.81C411.69,-525.66 412.64,-509.11 413.44,-495.18\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"416.93,-495.34 414.01,-485.15 409.95,-494.94 416.93,-495.34\"/>\n",
              "</g>\n",
              "<!-- 523449920 -->\n",
              "<g id=\"node523449920\" class=\"node\">\n",
              "<title>523449920</title>\n",
              "<polygon fill=\"#ffdeaf\" stroke=\"#ff8a4f\" points=\"461,-727 387,-727 387,-691 461,-691 461,-727\"/>\n",
              "<text text-anchor=\"start\" x=\"403.5\" y=\"-712.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">args[2]</text>\n",
              "<text text-anchor=\"start\" x=\"395\" y=\"-698.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,784]</text>\n",
              "</g>\n",
              "<!-- 523449920&#45;&gt;59282016 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>523449920&#45;&gt;59282016</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M420.96,-690.86C419.16,-680.71 416.85,-667.63 414.81,-656.12\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"418.22,-655.34 413.03,-646.11 411.33,-656.57 418.22,-655.34\"/>\n",
              "</g>\n",
              "<!-- 527083520 -->\n",
              "<g id=\"node527083520\" class=\"node\">\n",
              "<title>527083520</title>\n",
              "<polygon fill=\"#ffdeaf\" stroke=\"#ff8a4f\" points=\"369,-727 205,-727 205,-691 369,-691 369,-727\"/>\n",
              "<text text-anchor=\"start\" x=\"213\" y=\"-712.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">args[0]: mlp/~/linear_0/w</text>\n",
              "<text text-anchor=\"start\" x=\"251\" y=\"-698.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[784,300]</text>\n",
              "</g>\n",
              "<!-- 527083520&#45;&gt;59282016 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>527083520&#45;&gt;59282016</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M343.58,-690.85C348.61,-688.51 353.49,-685.9 358,-683 370.09,-675.22 381.43,-664.1 390.42,-653.97\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"393.27,-656.03 397.09,-646.14 387.94,-651.49 393.27,-656.03\"/>\n",
              "</g>\n",
              "<!-- 548696480 -->\n",
              "<g id=\"node548696480\" class=\"node\">\n",
              "<title>548696480</title>\n",
              "<polygon fill=\"#ffdeaf\" stroke=\"#ff8a4f\" points=\"187,-727 25,-727 25,-691 187,-691 187,-727\"/>\n",
              "<text text-anchor=\"start\" x=\"33\" y=\"-712.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">args[0]: mlp/~/linear_0/b</text>\n",
              "<text text-anchor=\"start\" x=\"82\" y=\"-698.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[300]</text>\n",
              "</g>\n",
              "<!-- 548696480&#45;&gt;69211888 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>548696480&#45;&gt;69211888</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M145.49,-690.86C173.04,-678.9 209.97,-662.87 239.34,-650.12\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"240.81,-653.3 248.59,-646.11 238.02,-646.88 240.81,-653.3\"/>\n",
              "</g>\n",
              "<!-- 70422144 -->\n",
              "<g id=\"node70422144\" class=\"node\">\n",
              "<title>70422144</title>\n",
              "<g id=\"a_node70422144\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"452.5,-396 367.5,-396 367.5,-360 452.5,-360 452.5,-396\"/>\n",
              "<text text-anchor=\"start\" x=\"375.5\" y=\"-381.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">dot_general</text>\n",
              "<text text-anchor=\"start\" x=\"381\" y=\"-367.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,100]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 67211248&#45;&gt;70422144 -->\n",
              "<g id=\"edge8\" class=\"edge\">\n",
              "<title>67211248&#45;&gt;70422144</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M414.01,-448.81C413.31,-436.66 412.36,-420.11 411.56,-406.18\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"415.05,-405.94 410.99,-396.15 408.07,-406.34 415.05,-405.94\"/>\n",
              "</g>\n",
              "<!-- 139737768045680 -->\n",
              "<g id=\"node8\" class=\"node\">\n",
              "<title>139737768045680</title>\n",
              "<polygon fill=\"#e6d6ff\" stroke=\"#a261ff\" points=\"211,-574 157,-574 157,-538 211,-538 211,-574\"/>\n",
              "<text text-anchor=\"start\" x=\"175\" y=\"-553.3\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">0.0</text>\n",
              "</g>\n",
              "<!-- 139737768045680&#45;&gt;67211248 -->\n",
              "<g id=\"edge7\" class=\"edge\">\n",
              "<title>139737768045680&#45;&gt;67211248</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M201.48,-537.82C205.61,-534.68 210.22,-531.82 215,-530 248.05,-517.39 342.72,-538.52 374,-522 385.79,-515.77 395.18,-504.48 401.97,-493.83\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"405.02,-495.55 407.07,-485.16 398.99,-492 405.02,-495.55\"/>\n",
              "</g>\n",
              "<!-- 50960960 -->\n",
              "<g id=\"node50960960\" class=\"node\">\n",
              "<title>50960960</title>\n",
              "<g id=\"a_node50960960\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"447,-324 373,-324 373,-288 447,-288 447,-324\"/>\n",
              "<text text-anchor=\"start\" x=\"399\" y=\"-309.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">add</text>\n",
              "<text text-anchor=\"start\" x=\"381\" y=\"-295.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,100]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 70422144&#45;&gt;50960960 -->\n",
              "<g id=\"edge11\" class=\"edge\">\n",
              "<title>70422144&#45;&gt;50960960</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M410,-359.7C410,-351.98 410,-342.71 410,-334.11\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"413.5,-334.1 410,-324.1 406.5,-334.1 413.5,-334.1\"/>\n",
              "</g>\n",
              "<!-- 69926688 -->\n",
              "<g id=\"node69926688\" class=\"node\">\n",
              "<title>69926688</title>\n",
              "<g id=\"a_node69926688\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"349,-396 227,-396 227,-360 349,-360 349,-396\"/>\n",
              "<text text-anchor=\"start\" x=\"235\" y=\"-381.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">broadcast_in_dim</text>\n",
              "<text text-anchor=\"start\" x=\"259\" y=\"-367.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,100]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 69926688&#45;&gt;50960960 -->\n",
              "<g id=\"edge12\" class=\"edge\">\n",
              "<title>69926688&#45;&gt;50960960</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M317.84,-359.88C333.94,-350.64 354,-339.13 371.29,-329.21\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"373.11,-332.2 380.04,-324.19 369.63,-326.13 373.11,-332.2\"/>\n",
              "</g>\n",
              "<!-- 68842336 -->\n",
              "<g id=\"node68842336\" class=\"node\">\n",
              "<title>68842336</title>\n",
              "<g id=\"a_node68842336\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"452,-235 378,-235 378,-199 452,-199 452,-235\"/>\n",
              "<text text-anchor=\"start\" x=\"402\" y=\"-220.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">max</text>\n",
              "<text text-anchor=\"start\" x=\"386\" y=\"-206.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,100]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 50960960&#45;&gt;68842336 -->\n",
              "<g id=\"edge13\" class=\"edge\">\n",
              "<title>50960960&#45;&gt;68842336</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M410.99,-287.81C411.69,-275.66 412.64,-259.11 413.44,-245.18\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"416.93,-245.34 414.01,-235.15 409.95,-244.94 416.93,-245.34\"/>\n",
              "</g>\n",
              "<!-- 653177280 -->\n",
              "<g id=\"node653177280\" class=\"node\">\n",
              "<title>653177280</title>\n",
              "<polygon fill=\"#ffdeaf\" stroke=\"#ff8a4f\" points=\"360,-485 196,-485 196,-449 360,-449 360,-485\"/>\n",
              "<text text-anchor=\"start\" x=\"204\" y=\"-470.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">args[0]: mlp/~/linear_1/w</text>\n",
              "<text text-anchor=\"start\" x=\"242\" y=\"-456.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[300,100]</text>\n",
              "</g>\n",
              "<!-- 653177280&#45;&gt;70422144 -->\n",
              "<g id=\"edge9\" class=\"edge\">\n",
              "<title>653177280&#45;&gt;70422144</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M328.3,-448.98C338.53,-444.51 348.93,-439.17 358,-433 369.8,-424.97 381.03,-413.9 390.02,-403.88\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"392.82,-406 396.71,-396.15 387.52,-401.42 392.82,-406\"/>\n",
              "</g>\n",
              "<!-- 642729120 -->\n",
              "<g id=\"node642729120\" class=\"node\">\n",
              "<title>642729120</title>\n",
              "<polygon fill=\"#ffdeaf\" stroke=\"#ff8a4f\" points=\"178,-485 16,-485 16,-449 178,-449 178,-485\"/>\n",
              "<text text-anchor=\"start\" x=\"24\" y=\"-470.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">args[0]: mlp/~/linear_1/b</text>\n",
              "<text text-anchor=\"start\" x=\"73\" y=\"-456.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[100]</text>\n",
              "</g>\n",
              "<!-- 642729120&#45;&gt;69926688 -->\n",
              "<g id=\"edge10\" class=\"edge\">\n",
              "<title>642729120&#45;&gt;69926688</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M134.74,-448.81C165.31,-434.88 208.54,-415.19 241.3,-400.27\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"243.02,-403.33 250.67,-396 240.12,-396.96 243.02,-403.33\"/>\n",
              "</g>\n",
              "<!-- 71785552 -->\n",
              "<g id=\"node71785552\" class=\"node\">\n",
              "<title>71785552</title>\n",
              "<g id=\"a_node71785552\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"320.5,-146 235.5,-146 235.5,-110 320.5,-110 320.5,-146\"/>\n",
              "<text text-anchor=\"start\" x=\"243.5\" y=\"-131.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">dot_general</text>\n",
              "<text text-anchor=\"start\" x=\"252\" y=\"-117.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,10]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 68842336&#45;&gt;71785552 -->\n",
              "<g id=\"edge15\" class=\"edge\">\n",
              "<title>68842336&#45;&gt;71785552</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M387.93,-198.81C366.64,-185.29 336.78,-166.33 313.56,-151.58\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"315.33,-148.56 305.01,-146.15 311.58,-154.47 315.33,-148.56\"/>\n",
              "</g>\n",
              "<!-- 139737768255984 -->\n",
              "<g id=\"node15\" class=\"node\">\n",
              "<title>139737768255984</title>\n",
              "<polygon fill=\"#e6d6ff\" stroke=\"#a261ff\" points=\"211,-324 157,-324 157,-288 211,-288 211,-324\"/>\n",
              "<text text-anchor=\"start\" x=\"175\" y=\"-303.3\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">0.0</text>\n",
              "</g>\n",
              "<!-- 139737768255984&#45;&gt;68842336 -->\n",
              "<g id=\"edge14\" class=\"edge\">\n",
              "<title>139737768255984&#45;&gt;68842336</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M201.48,-287.82C205.61,-284.68 210.22,-281.82 215,-280 248.05,-267.39 342.72,-288.52 374,-272 385.79,-265.77 395.18,-254.48 401.97,-243.83\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"405.02,-245.55 407.07,-235.16 398.99,-242 405.02,-245.55\"/>\n",
              "</g>\n",
              "<!-- 72015984 -->\n",
              "<g id=\"node72015984\" class=\"node\">\n",
              "<title>72015984</title>\n",
              "<g id=\"a_node72015984\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#aabfff\" stroke=\"#0053d6\" stroke-width=\"2\" points=\"251,-74 183,-74 183,-24 251,-24 251,-74\"/>\n",
              "<text text-anchor=\"start\" x=\"197.5\" y=\"-59.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">output</text>\n",
              "<text text-anchor=\"start\" x=\"206\" y=\"-45.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">add</text>\n",
              "<text text-anchor=\"start\" x=\"191\" y=\"-31.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,10]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 71785552&#45;&gt;72015984 -->\n",
              "<g id=\"edge18\" class=\"edge\">\n",
              "<title>71785552&#45;&gt;72015984</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M264.46,-109.91C257.94,-101.67 249.89,-91.52 242.29,-81.92\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"245.01,-79.71 236.05,-74.05 239.52,-84.06 245.01,-79.71\"/>\n",
              "</g>\n",
              "<!-- 70466816 -->\n",
              "<g id=\"node70466816\" class=\"node\">\n",
              "<title>70466816</title>\n",
              "<g id=\"a_node70466816\"><a xlink:title=\" \">\n",
              "<polygon fill=\"#fff26e\" stroke=\"#ffdb13\" points=\"217,-146 95,-146 95,-110 217,-110 217,-146\"/>\n",
              "<text text-anchor=\"start\" x=\"103\" y=\"-131.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">broadcast_in_dim</text>\n",
              "<text text-anchor=\"start\" x=\"130\" y=\"-117.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[8,10]</text>\n",
              "</a>\n",
              "</g>\n",
              "</g>\n",
              "<!-- 70466816&#45;&gt;72015984 -->\n",
              "<g id=\"edge19\" class=\"edge\">\n",
              "<title>70466816&#45;&gt;72015984</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M169.54,-109.91C176.06,-101.67 184.11,-91.52 191.71,-81.92\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"194.48,-84.06 197.95,-74.05 188.99,-79.71 194.48,-84.06\"/>\n",
              "</g>\n",
              "<!-- 689771200 -->\n",
              "<g id=\"node689771200\" class=\"node\">\n",
              "<title>689771200</title>\n",
              "<polygon fill=\"#ffdeaf\" stroke=\"#ff8a4f\" points=\"360,-235 196,-235 196,-199 360,-199 360,-235\"/>\n",
              "<text text-anchor=\"start\" x=\"204\" y=\"-220.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">args[0]: mlp/~/linear_2/w</text>\n",
              "<text text-anchor=\"start\" x=\"245.5\" y=\"-206.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[100,10]</text>\n",
              "</g>\n",
              "<!-- 689771200&#45;&gt;71785552 -->\n",
              "<g id=\"edge16\" class=\"edge\">\n",
              "<title>689771200&#45;&gt;71785552</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M278,-198.81C278,-186.66 278,-170.11 278,-156.18\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"281.5,-156.15 278,-146.15 274.5,-156.15 281.5,-156.15\"/>\n",
              "</g>\n",
              "<!-- 690080640 -->\n",
              "<g id=\"node690080640\" class=\"node\">\n",
              "<title>690080640</title>\n",
              "<polygon fill=\"#ffdeaf\" stroke=\"#ff8a4f\" points=\"178,-235 16,-235 16,-199 178,-199 178,-235\"/>\n",
              "<text text-anchor=\"start\" x=\"24\" y=\"-220.8\" font-family=\"Times,serif\" font-weight=\"bold\" font-size=\"14.00\">args[0]: mlp/~/linear_2/b</text>\n",
              "<text text-anchor=\"start\" x=\"76.5\" y=\"-206.8\" font-family=\"Times,serif\" font-size=\"14.00\">f32[10]</text>\n",
              "</g>\n",
              "<!-- 690080640&#45;&gt;70466816 -->\n",
              "<g id=\"edge17\" class=\"edge\">\n",
              "<title>690080640&#45;&gt;70466816</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M108.66,-198.81C117.22,-186.18 129,-168.81 138.67,-154.55\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"141.65,-156.39 144.37,-146.15 135.86,-152.46 141.65,-156.39\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x7f173da873d0>"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "def f(x):\n",
        "  return hk.nets.MLP([300, 100, 10])(x)\n",
        "\n",
        "f = hk.transform(f)\n",
        "\n",
        "rng = jax.random.PRNGKey(42)\n",
        "x = jnp.ones([8, 28 * 28])\n",
        "params = f.init(rng, x)\n",
        "\n",
        "dot = hk.to_dot(f.apply)(params, None, x)\n",
        "graphviz.Source(dot)"
      ]
    }
  ],
  "metadata": {
    "colab": {}
  },
  "nbformat": 4,
  "nbformat_minor": 0
}