tutorial/source/model_rendering.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "857effec",
   "metadata": {},
   "source": [
    "# Automatic rendering of Pyro models\n",
    "\n",
    "In this tutorial we will demonstrate how to create beautiful visualizations of your probabilistic graphical models using [pyro.render_model()](https://docs.pyro.ai/en/latest/infer.util.html#pyro.infer.inspect.render_model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8a068eb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "import pyro.distributions.constraints as constraints\n",
    "\n",
    "smoke_test = ('CI' in os.environ)\n",
    "assert pyro.__version__.startswith('1.9.1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1266c17",
   "metadata": {},
   "source": [
    "## A Simple Example\n",
    "\n",
    "The visualization interface can be readily used with your models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "855a7d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(data):\n",
    "    m = pyro.sample(\"m\", dist.Normal(0, 1))\n",
    "    sd = pyro.sample(\"sd\", dist.LogNormal(m, 1))\n",
    "    with pyro.plate(\"N\", len(data)):\n",
    "        pyro.sample(\"obs\", dist.Normal(m, sd), obs=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e1e9628e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\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 8.1.0 (0)\n -->\n<!-- Pages: 1 -->\n<svg width=\"107pt\" height=\"229pt\"\n viewBox=\"0.00 0.00 107.00 229.25\" 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 225.25)\">\n<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-225.25 103,-225.25 103,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster_N</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"17,-8 17,-85.25 91,-85.25 91,-8 17,-8\"/>\n<text text-anchor=\"middle\" x=\"77\" y=\"-15.95\" font-family=\"Times,serif\" font-size=\"14.00\">N</text>\n</g>\n<!-- m -->\n<g id=\"node1\" class=\"node\">\n<title>m</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"54\" cy=\"-203.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"54\" y=\"-198.57\" font-family=\"Times,serif\" font-size=\"14.00\">m</text>\n</g>\n<!-- sd -->\n<g id=\"node2\" class=\"node\">\n<title>sd</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"27\" cy=\"-131.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"27\" y=\"-126.58\" font-family=\"Times,serif\" font-size=\"14.00\">sd</text>\n</g>\n<!-- m&#45;&gt;sd -->\n<g id=\"edge1\" class=\"edge\">\n<title>m&#45;&gt;sd</title>\n<path fill=\"none\" stroke=\"black\" d=\"M47.46,-185.3C44.45,-177.48 40.79,-168 37.41,-159.23\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"40.34,-158.11 33.48,-150.04 33.81,-160.63 40.34,-158.11\"/>\n</g>\n<!-- obs -->\n<g id=\"node3\" class=\"node\">\n<title>obs</title>\n<ellipse fill=\"gray\" stroke=\"black\" cx=\"54\" cy=\"-59.25\" rx=\"28.59\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"54\" y=\"-54.58\" font-family=\"Times,serif\" font-size=\"14.00\">obs</text>\n</g>\n<!-- m&#45;&gt;obs -->\n<g id=\"edge2\" class=\"edge\">\n<title>m&#45;&gt;obs</title>\n<path fill=\"none\" stroke=\"black\" d=\"M57.65,-185.16C59.68,-174.82 61.98,-161.34 63,-149.25 64.34,-133.31 64.34,-129.19 63,-113.25 62.31,-105.03 61.02,-96.17 59.63,-88.1\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"62.91,-87.54 57.65,-78.34 56.02,-88.81 62.91,-87.54\"/>\n</g>\n<!-- sd&#45;&gt;obs -->\n<g id=\"edge3\" class=\"edge\">\n<title>sd&#45;&gt;obs</title>\n<path fill=\"none\" stroke=\"black\" d=\"M33.54,-113.3C36.55,-105.48 40.21,-96 43.59,-87.23\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"47.19,-88.63 47.52,-78.04 40.66,-86.11 47.19,-88.63\"/>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f06e8976510>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = torch.ones(10)\n",
    "pyro.render_model(model, model_args=(data,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb1f65be",
   "metadata": {},
   "source": [
    "The visualization can be saved to a file by providing `filename='path'` to `pyro.render_model`. You can use different formats such as PDF or PNG by changing the filename's suffix.\n",
    "When not saving to a file (`filename=None`), you can also change the format with `graph.format = 'pdf'` where `graph` is the object returned by `pyro.render_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "700a0917",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = pyro.render_model(model, model_args=(data,), filename=\"model.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a4ea614",
   "metadata": {},
   "source": [
    "## Tweaking the visualization\n",
    "\n",
    "As `pyro.render_model` returns an object of type `graphviz.dot.Digraph`, you can further improve the visualization of this graph.\n",
    "For example, you could use the [unflatten preprocessor](https://graphviz.readthedocs.io/en/stable/api.html#graphviz.unflatten) to improve the layout aspect ratio for more complex models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "01a4b74b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mace(positions, annotations):\n",
    "    \"\"\"\n",
    "    This model corresponds to the plate diagram in Figure 3 of https://www.aclweb.org/anthology/Q18-1040.pdf.\n",
    "    \"\"\"\n",
    "    num_annotators = int(torch.max(positions)) + 1\n",
    "    num_classes = int(torch.max(annotations)) + 1\n",
    "    num_items, num_positions = annotations.shape\n",
    "\n",
    "    with pyro.plate(\"annotator\", num_annotators):\n",
    "        epsilon = pyro.sample(\"ε\", dist.Dirichlet(torch.full((num_classes,), 10.)))\n",
    "        theta = pyro.sample(\"θ\", dist.Beta(0.5, 0.5))\n",
    "\n",
    "    with pyro.plate(\"item\", num_items, dim=-2):\n",
    "        # NB: using constant logits for discrete uniform prior\n",
    "        # (NumPyro does not have DiscreteUniform distribution yet)\n",
    "        c = pyro.sample(\"c\", dist.Categorical(logits=torch.zeros(num_classes)))\n",
    "\n",
    "        with pyro.plate(\"position\", num_positions):\n",
    "            s = pyro.sample(\"s\", dist.Bernoulli(1 - theta[positions]))\n",
    "            probs = torch.where(\n",
    "                s[..., None] == 0, F.one_hot(c, num_classes).float(), epsilon[positions]\n",
    "            )\n",
    "            pyro.sample(\"y\", dist.Categorical(probs), obs=annotations)\n",
    "\n",
    "\n",
    "positions = torch.tensor([1, 1, 1, 2, 3, 4, 5])\n",
    "# fmt: off\n",
    "annotations = torch.tensor([\n",
    "    [1, 3, 1, 2, 2, 2, 1, 3, 2, 2, 4, 2, 1, 2, 1,\n",
    "     1, 1, 1, 2, 2, 2, 2, 2, 2, 1, 1, 2, 1, 1, 1,\n",
    "     1, 3, 1, 2, 2, 4, 2, 2, 3, 1, 1, 1, 2, 1, 2],\n",
    "    [1, 3, 1, 2, 2, 2, 2, 3, 2, 3, 4, 2, 1, 2, 2,\n",
    "     1, 1, 1, 2, 2, 2, 2, 2, 2, 1, 1, 3, 1, 1, 1,\n",
    "     1, 3, 1, 2, 2, 3, 2, 3, 3, 1, 1, 2, 3, 2, 2],\n",
    "    [1, 3, 2, 2, 2, 2, 2, 3, 2, 2, 4, 2, 1, 2, 1,\n",
    "     1, 1, 1, 2, 2, 2, 2, 2, 1, 1, 1, 2, 1, 1, 2,\n",
    "     1, 3, 1, 2, 2, 3, 1, 2, 3, 1, 1, 1, 2, 1, 2],\n",
    "    [1, 4, 2, 3, 3, 3, 2, 3, 2, 2, 4, 3, 1, 3, 1,\n",
    "     2, 1, 1, 2, 1, 2, 2, 3, 2, 1, 1, 2, 1, 1, 1,\n",
    "     1, 3, 1, 2, 3, 4, 2, 3, 3, 1, 1, 2, 2, 1, 2],\n",
    "    [1, 3, 1, 1, 2, 3, 1, 4, 2, 2, 4, 3, 1, 2, 1,\n",
    "     1, 1, 1, 2, 3, 2, 2, 2, 2, 1, 1, 2, 1, 1, 1,\n",
    "     1, 2, 1, 2, 2, 3, 2, 2, 4, 1, 1, 1, 2, 1, 2],\n",
    "    [1, 3, 2, 2, 2, 2, 1, 3, 2, 2, 4, 4, 1, 1, 1,\n",
    "     1, 1, 1, 2, 2, 2, 2, 2, 2, 1, 1, 2, 1, 1, 2,\n",
    "     1, 3, 1, 2, 3, 4, 3, 3, 3, 1, 1, 1, 2, 1, 2],\n",
    "    [1, 4, 2, 1, 2, 2, 1, 3, 3, 3, 4, 3, 1, 2, 1,\n",
    "     1, 1, 1, 1, 2, 2, 1, 2, 2, 1, 1, 2, 1, 1, 1,\n",
    "     1, 3, 1, 2, 2, 3, 2, 3, 2, 1, 1, 1, 2, 1, 2],\n",
    "]).T\n",
    "# fmt: on\n",
    "\n",
    "# we subtract 1 because the first index starts with 0 in Python\n",
    "positions -= 1\n",
    "annotations -= 1\n",
    "\n",
    "mace_graph = pyro.render_model(mace, model_args=(positions, annotations))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "311896af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\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 8.1.0 (0)\n -->\n<!-- Pages: 1 -->\n<svg width=\"202pt\" height=\"300pt\"\n viewBox=\"0.00 0.00 202.00 299.75\" 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 295.75)\">\n<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-295.75 198,-295.75 198,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster_annotator</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"8,-206.5 8,-283.75 150,-283.75 150,-206.5 8,-206.5\"/>\n<text text-anchor=\"middle\" x=\"107.12\" y=\"-214.45\" font-family=\"Times,serif\" font-size=\"14.00\">annotator</text>\n</g>\n<g id=\"clust2\" class=\"cluster\">\n<title>cluster_item</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"35,-8 35,-198.5 186,-198.5 186,-8 35,-8\"/>\n<text text-anchor=\"middle\" x=\"161.88\" y=\"-15.95\" font-family=\"Times,serif\" font-size=\"14.00\">item</text>\n</g>\n<g id=\"clust3\" class=\"cluster\">\n<title>cluster_position</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"105,-41.25 105,-190.5 178,-190.5 178,-41.25 105,-41.25\"/>\n<text text-anchor=\"middle\" x=\"141.5\" y=\"-49.2\" font-family=\"Times,serif\" font-size=\"14.00\">position</text>\n</g>\n<!-- ε -->\n<g id=\"node1\" class=\"node\">\n<title>ε</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"43\" cy=\"-257.75\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"43\" y=\"-253.07\" font-family=\"Times,serif\" font-size=\"14.00\">ε</text>\n</g>\n<!-- y -->\n<g id=\"node5\" class=\"node\">\n<title>y</title>\n<ellipse fill=\"gray\" stroke=\"black\" cx=\"140\" cy=\"-92.5\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"140\" y=\"-87.83\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n</g>\n<!-- ε&#45;&gt;y -->\n<g id=\"edge2\" class=\"edge\">\n<title>ε&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M35.27,-240.37C25.72,-217.51 12.82,-175.55 31,-146.5 46.75,-121.33 78.26,-107.8 103.36,-100.73\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"104.2,-103.88 113.02,-98 102.46,-97.09 104.2,-103.88\"/>\n</g>\n<!-- θ -->\n<g id=\"node2\" class=\"node\">\n<title>θ</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"115\" cy=\"-257.75\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"115\" y=\"-253.07\" font-family=\"Times,serif\" font-size=\"14.00\">θ</text>\n</g>\n<!-- s -->\n<g id=\"node4\" class=\"node\">\n<title>s</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"142\" cy=\"-164.5\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"142\" y=\"-159.82\" font-family=\"Times,serif\" font-size=\"14.00\">s</text>\n</g>\n<!-- θ&#45;&gt;s -->\n<g id=\"edge1\" class=\"edge\">\n<title>θ&#45;&gt;s</title>\n<path fill=\"none\" stroke=\"black\" d=\"M120.08,-239.59C123.96,-226.46 129.4,-208.1 133.88,-192.94\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"137.46,-194.17 136.94,-183.59 130.75,-192.18 137.46,-194.17\"/>\n</g>\n<!-- c -->\n<g id=\"node3\" class=\"node\">\n<title>c</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"70\" cy=\"-164.5\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"70\" y=\"-159.82\" font-family=\"Times,serif\" font-size=\"14.00\">c</text>\n</g>\n<!-- c&#45;&gt;y -->\n<g id=\"edge3\" class=\"edge\">\n<title>c&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M84.5,-149C94.17,-139.33 107.06,-126.44 117.98,-115.52\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"119.92,-118.53 124.51,-108.99 114.97,-113.58 119.92,-118.53\"/>\n</g>\n<!-- s&#45;&gt;y -->\n<g id=\"edge4\" class=\"edge\">\n<title>s&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M141.51,-146.2C141.29,-138.74 141.04,-129.82 140.8,-121.47\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"144.27,-121.5 140.49,-111.6 137.28,-121.7 144.27,-121.5\"/>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f06e8ad44d0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# default layout\n",
    "mace_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8babebb4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\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 8.1.0 (0)\n -->\n<!-- Pages: 1 -->\n<svg width=\"272pt\" height=\"292pt\"\n viewBox=\"0.00 0.00 272.00 291.75\" 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 287.75)\">\n<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-287.75 268,-287.75 268,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster_annotator</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"8,-198.5 8,-275.75 150,-275.75 150,-198.5 8,-198.5\"/>\n<text text-anchor=\"middle\" x=\"107.12\" y=\"-206.45\" font-family=\"Times,serif\" font-size=\"14.00\">annotator</text>\n</g>\n<g id=\"clust2\" class=\"cluster\">\n<title>cluster_item</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"158,-8 158,-275.75 256,-275.75 256,-8 158,-8\"/>\n<text text-anchor=\"middle\" x=\"231.88\" y=\"-15.95\" font-family=\"Times,serif\" font-size=\"14.00\">item</text>\n</g>\n<g id=\"clust3\" class=\"cluster\">\n<title>cluster_position</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"166,-41.25 166,-190.5 239,-190.5 239,-41.25 166,-41.25\"/>\n<text text-anchor=\"middle\" x=\"202.5\" y=\"-49.2\" font-family=\"Times,serif\" font-size=\"14.00\">position</text>\n</g>\n<!-- ε -->\n<g id=\"node1\" class=\"node\">\n<title>ε</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"43\" cy=\"-249.75\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"43\" y=\"-245.07\" font-family=\"Times,serif\" font-size=\"14.00\">ε</text>\n</g>\n<!-- y -->\n<g id=\"node4\" class=\"node\">\n<title>y</title>\n<ellipse fill=\"gray\" stroke=\"black\" cx=\"201\" cy=\"-92.5\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"201\" y=\"-87.83\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n</g>\n<!-- ε&#45;&gt;y -->\n<g id=\"edge1\" class=\"edge\">\n<title>ε&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M53.12,-232.62C59.89,-222.38 69.29,-209.11 79,-198.5 109.03,-165.67 148.9,-133.06 174.75,-113.09\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"176.46,-115.44 182.27,-106.58 172.2,-109.87 176.46,-115.44\"/>\n</g>\n<!-- θ -->\n<g id=\"node2\" class=\"node\">\n<title>θ</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"115\" cy=\"-249.75\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"115\" y=\"-245.07\" font-family=\"Times,serif\" font-size=\"14.00\">θ</text>\n</g>\n<!-- s -->\n<g id=\"node3\" class=\"node\">\n<title>s</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"201\" cy=\"-164.5\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"201\" y=\"-159.82\" font-family=\"Times,serif\" font-size=\"14.00\">s</text>\n</g>\n<!-- θ&#45;&gt;s -->\n<g id=\"edge2\" class=\"edge\">\n<title>θ&#45;&gt;s</title>\n<path fill=\"none\" stroke=\"black\" d=\"M125.51,-232.81C132.8,-222.4 143.13,-208.86 154,-198.5 159.76,-193.01 166.49,-187.74 173.03,-183.09\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"174.43,-185.68 180.72,-177.15 170.48,-179.9 174.43,-185.68\"/>\n</g>\n<!-- s&#45;&gt;y -->\n<g id=\"edge4\" class=\"edge\">\n<title>s&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M201,-146.2C201,-138.74 201,-129.82 201,-121.47\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"204.5,-121.6 201,-111.6 197.5,-121.6 204.5,-121.6\"/>\n</g>\n<!-- c -->\n<g id=\"node5\" class=\"node\">\n<title>c</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"221\" cy=\"-249.75\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"221\" y=\"-245.07\" font-family=\"Times,serif\" font-size=\"14.00\">c</text>\n</g>\n<!-- c&#45;&gt;y -->\n<g id=\"edge3\" class=\"edge\">\n<title>c&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M228.11,-232.13C235.86,-211.62 246.03,-175.87 237,-146.5 233.78,-136.04 227.72,-125.79 221.47,-117.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"223.63,-115.2 214.75,-109.43 218.1,-119.48 223.63,-115.2\"/>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.sources.Source at 0x7f07a2e06150>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# layout after processing the layout with unflatten\n",
    "mace_graph.unflatten(stagger=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50a92902",
   "metadata": {},
   "source": [
    "## Rendering the parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74f32e20",
   "metadata": {},
   "source": [
    "We can render the parameters defined as `pyro.param` by setting `render_params=True` in `pyro.render_model`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "645df936",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(data):\n",
    "    sigma = pyro.param(\"sigma\", torch.tensor([1.]), constraint=constraints.positive)\n",
    "    mu = pyro.param(\"mu\", torch.tensor([0.]))\n",
    "    x = pyro.sample(\"x\", dist.Normal(mu, sigma))\n",
    "    y = pyro.sample(\"y\", dist.LogNormal(x, 1))\n",
    "    with pyro.plate(\"N\", len(data)):\n",
    "        pyro.sample(\"z\", dist.Normal(x, y), obs=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "66fc9f55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\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 8.1.0 (0)\n -->\n<!-- Pages: 1 -->\n<svg width=\"106pt\" height=\"283pt\"\n viewBox=\"0.00 0.00 106.00 282.50\" 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 278.5)\">\n<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-278.5 102,-278.5 102,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster_N</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"20,-8 20,-85.25 90,-85.25 90,-8 20,-8\"/>\n<text text-anchor=\"middle\" x=\"76\" y=\"-15.95\" font-family=\"Times,serif\" font-size=\"14.00\">N</text>\n</g>\n<!-- x -->\n<g id=\"node1\" class=\"node\">\n<title>x</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"54\" cy=\"-203.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"54\" y=\"-198.57\" font-family=\"Times,serif\" font-size=\"14.00\">x</text>\n</g>\n<!-- y -->\n<g id=\"node2\" class=\"node\">\n<title>y</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"27\" cy=\"-131.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"27\" y=\"-126.58\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n</g>\n<!-- x&#45;&gt;y -->\n<g id=\"edge1\" class=\"edge\">\n<title>x&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M47.46,-185.3C44.45,-177.48 40.79,-168 37.41,-159.23\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"40.34,-158.11 33.48,-150.04 33.81,-160.63 40.34,-158.11\"/>\n</g>\n<!-- z -->\n<g id=\"node5\" class=\"node\">\n<title>z</title>\n<ellipse fill=\"gray\" stroke=\"black\" cx=\"55\" cy=\"-59.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"55\" y=\"-54.58\" font-family=\"Times,serif\" font-size=\"14.00\">z</text>\n</g>\n<!-- x&#45;&gt;z -->\n<g id=\"edge2\" class=\"edge\">\n<title>x&#45;&gt;z</title>\n<path fill=\"none\" stroke=\"black\" d=\"M57.65,-185.16C59.68,-174.82 61.98,-161.34 63,-149.25 64.34,-133.31 64.2,-129.21 63,-113.25 62.38,-105.05 61.24,-96.2 60.01,-88.12\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"63.32,-87.67 58.25,-78.36 56.41,-88.8 63.32,-87.67\"/>\n</g>\n<!-- y&#45;&gt;z -->\n<g id=\"edge3\" class=\"edge\">\n<title>y&#45;&gt;z</title>\n<path fill=\"none\" stroke=\"black\" d=\"M33.64,-113.66C36.84,-105.64 40.77,-95.82 44.38,-86.79\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"47.95,-88.3 48.41,-77.72 41.45,-85.7 47.95,-88.3\"/>\n</g>\n<!-- sigma -->\n<g id=\"node3\" class=\"node\">\n<title>sigma</title>\n<polygon fill=\"white\" stroke=\"none\" points=\"50.38,-274.5 7.62,-274.5 7.62,-257.25 50.38,-257.25 50.38,-274.5\"/>\n<text text-anchor=\"middle\" x=\"29\" y=\"-261.2\" font-family=\"Times,serif\" font-size=\"14.00\">sigma</text>\n</g>\n<!-- sigma&#45;&gt;x -->\n<g id=\"edge4\" class=\"edge\">\n<title>sigma&#45;&gt;x</title>\n<path fill=\"none\" stroke=\"black\" d=\"M32.07,-257.44C34.8,-250.8 38.99,-240.66 42.97,-231\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"46.55,-232.49 47.13,-221.91 40.08,-229.82 46.55,-232.49\"/>\n</g>\n<!-- mu -->\n<g id=\"node4\" class=\"node\">\n<title>mu</title>\n<polygon fill=\"white\" stroke=\"none\" points=\"91.25,-274.5 68.75,-274.5 68.75,-257.25 91.25,-257.25 91.25,-274.5\"/>\n<text text-anchor=\"middle\" x=\"80\" y=\"-261.2\" font-family=\"Times,serif\" font-size=\"14.00\">mu</text>\n</g>\n<!-- mu&#45;&gt;x -->\n<g id=\"edge5\" class=\"edge\">\n<title>mu&#45;&gt;x</title>\n<path fill=\"none\" stroke=\"black\" d=\"M76.81,-257.44C73.96,-250.8 69.61,-240.66 65.47,-231\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"68.3,-229.72 61.14,-221.91 61.87,-232.48 68.3,-229.72\"/>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f07bc1d1910>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = torch.ones(10)\n",
    "pyro.render_model(model, model_args=(data,), render_params=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f09e73e1",
   "metadata": {},
   "source": [
    "## Distribution and Constraint annotations\n",
    "\n",
    "It is possible to display the distribution of each RV in the generated plot by providing `render_distributions=True` when calling `pyro.render_model`. The constraints associated with parameters are also displayed when `render_distributions=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8130359c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\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 8.1.0 (0)\n -->\n<!-- Pages: 1 -->\n<svg width=\"416pt\" height=\"360pt\"\n viewBox=\"0.00 0.00 416.38 359.50\" 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 355.5)\">\n<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-355.5 412.38,-355.5 412.38,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster_N</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"20,-8 20,-85.25 90,-85.25 90,-8 20,-8\"/>\n<text text-anchor=\"middle\" x=\"76\" y=\"-15.95\" font-family=\"Times,serif\" font-size=\"14.00\">N</text>\n</g>\n<!-- x -->\n<g id=\"node1\" class=\"node\">\n<title>x</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"54\" cy=\"-203.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"54\" y=\"-198.57\" font-family=\"Times,serif\" font-size=\"14.00\">x</text>\n</g>\n<!-- y -->\n<g id=\"node2\" class=\"node\">\n<title>y</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"27\" cy=\"-131.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"27\" y=\"-126.58\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n</g>\n<!-- x&#45;&gt;y -->\n<g id=\"edge1\" class=\"edge\">\n<title>x&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M47.46,-185.3C44.45,-177.48 40.79,-168 37.41,-159.23\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"40.34,-158.11 33.48,-150.04 33.81,-160.63 40.34,-158.11\"/>\n</g>\n<!-- z -->\n<g id=\"node5\" class=\"node\">\n<title>z</title>\n<ellipse fill=\"gray\" stroke=\"black\" cx=\"55\" cy=\"-59.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"55\" y=\"-54.58\" font-family=\"Times,serif\" font-size=\"14.00\">z</text>\n</g>\n<!-- x&#45;&gt;z -->\n<g id=\"edge2\" class=\"edge\">\n<title>x&#45;&gt;z</title>\n<path fill=\"none\" stroke=\"black\" d=\"M57.65,-185.16C59.68,-174.82 61.98,-161.34 63,-149.25 64.34,-133.31 64.2,-129.21 63,-113.25 62.38,-105.05 61.24,-96.2 60.01,-88.12\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"63.32,-87.67 58.25,-78.36 56.41,-88.8 63.32,-87.67\"/>\n</g>\n<!-- y&#45;&gt;z -->\n<g id=\"edge3\" class=\"edge\">\n<title>y&#45;&gt;z</title>\n<path fill=\"none\" stroke=\"black\" d=\"M33.64,-113.66C36.84,-105.64 40.77,-95.82 44.38,-86.79\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"47.95,-88.3 48.41,-77.72 41.45,-85.7 47.95,-88.3\"/>\n</g>\n<!-- sigma -->\n<g id=\"node3\" class=\"node\">\n<title>sigma</title>\n<polygon fill=\"white\" stroke=\"none\" points=\"50.38,-313 7.62,-313 7.62,-295.75 50.38,-295.75 50.38,-313\"/>\n<text text-anchor=\"middle\" x=\"29\" y=\"-299.7\" font-family=\"Times,serif\" font-size=\"14.00\">sigma</text>\n</g>\n<!-- sigma&#45;&gt;x -->\n<g id=\"edge4\" class=\"edge\">\n<title>sigma&#45;&gt;x</title>\n<path fill=\"none\" stroke=\"black\" d=\"M30.88,-295.92C34.26,-282.51 41.54,-253.64 47.04,-231.87\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"50.66,-232.79 49.71,-222.24 43.88,-231.08 50.66,-232.79\"/>\n</g>\n<!-- mu -->\n<g id=\"node4\" class=\"node\">\n<title>mu</title>\n<polygon fill=\"white\" stroke=\"none\" points=\"91.25,-313 68.75,-313 68.75,-295.75 91.25,-295.75 91.25,-313\"/>\n<text text-anchor=\"middle\" x=\"80\" y=\"-299.7\" font-family=\"Times,serif\" font-size=\"14.00\">mu</text>\n</g>\n<!-- mu&#45;&gt;x -->\n<g id=\"edge5\" class=\"edge\">\n<title>mu&#45;&gt;x</title>\n<path fill=\"none\" stroke=\"black\" d=\"M78.05,-295.92C74.53,-282.51 66.96,-253.64 61.24,-231.87\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"64.38,-231.03 58.46,-222.24 57.61,-232.8 64.38,-231.03\"/>\n</g>\n<!-- distribution_description_node -->\n<g id=\"node6\" class=\"node\">\n<title>distribution_description_node</title>\n<text text-anchor=\"start\" x=\"117.62\" y=\"-334.2\" font-family=\"Times,serif\" font-size=\"14.00\">x ~ Normal</text>\n<text text-anchor=\"start\" x=\"117.62\" y=\"-316.95\" font-family=\"Times,serif\" font-size=\"14.00\">y ~ LogNormal</text>\n<text text-anchor=\"start\" x=\"117.62\" y=\"-299.7\" font-family=\"Times,serif\" font-size=\"14.00\">z ~ Normal</text>\n<text text-anchor=\"start\" x=\"117.62\" y=\"-282.45\" font-family=\"Times,serif\" font-size=\"14.00\">sigma : GreaterThan(lower_bound=0.0)</text>\n<text text-anchor=\"start\" x=\"117.62\" y=\"-265.2\" font-family=\"Times,serif\" font-size=\"14.00\">mu : Real()</text>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f06e8a04710>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = torch.ones(10)\n",
    "pyro.render_model(model, model_args=(data,), render_params=True ,render_distributions=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a1e4c4e",
   "metadata": {},
   "source": [
    "In the above plot **'~'** denotes the distribution of RV and **'$\\in$'** denotes the constraint of parameter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dbe4925",
   "metadata": {},
   "source": [
    "## Overlapping non-nested plates\n",
    "\n",
    "Note that overlapping non-nested plates may be drawn as multiple rectangles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9f1a4ebf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model():\n",
    "    plate1 = pyro.plate(\"plate1\", 2, dim=-2)\n",
    "    plate2 = pyro.plate(\"plate2\", 3, dim=-1)\n",
    "    with plate1:\n",
    "        x = pyro.sample(\"x\", dist.Normal(0, 1))\n",
    "    with plate1, plate2:\n",
    "        y = pyro.sample(\"y\", dist.Normal(x, 1))\n",
    "    with plate2:\n",
    "        pyro.sample(\"z\", dist.Normal(y.sum(-2, True), 1), obs=torch.zeros(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8514d7df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\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 8.1.0 (0)\n -->\n<!-- Pages: 1 -->\n<svg width=\"110pt\" height=\"292pt\"\n viewBox=\"0.00 0.00 110.00 291.75\" 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 287.75)\">\n<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-287.75 106,-287.75 106,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster_plate1</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"8,-93.25 8,-275.75 94,-275.75 94,-93.25 8,-93.25\"/>\n<text text-anchor=\"middle\" x=\"63.5\" y=\"-101.2\" font-family=\"Times,serif\" font-size=\"14.00\">plate1</text>\n</g>\n<g id=\"clust2\" class=\"cluster\">\n<title>cluster_plate2</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"16,-126.5 16,-203.75 86,-203.75 86,-126.5 16,-126.5\"/>\n<text text-anchor=\"middle\" x=\"55.5\" y=\"-134.45\" font-family=\"Times,serif\" font-size=\"14.00\">plate2</text>\n</g>\n<g id=\"clust3\" class=\"cluster\">\n<title>cluster_plate2__CLONE</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"16,-8 16,-85.25 86,-85.25 86,-8 16,-8\"/>\n<text text-anchor=\"middle\" x=\"55.5\" y=\"-15.95\" font-family=\"Times,serif\" font-size=\"14.00\">plate2</text>\n</g>\n<!-- x -->\n<g id=\"node1\" class=\"node\">\n<title>x</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"51\" cy=\"-249.75\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"51\" y=\"-245.07\" font-family=\"Times,serif\" font-size=\"14.00\">x</text>\n</g>\n<!-- y -->\n<g id=\"node2\" class=\"node\">\n<title>y</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"51\" cy=\"-177.75\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"51\" y=\"-173.07\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n</g>\n<!-- x&#45;&gt;y -->\n<g id=\"edge1\" class=\"edge\">\n<title>x&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M51,-231.45C51,-223.99 51,-215.07 51,-206.72\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54.5,-206.85 51,-196.85 47.5,-206.85 54.5,-206.85\"/>\n</g>\n<!-- z -->\n<g id=\"node3\" class=\"node\">\n<title>z</title>\n<ellipse fill=\"gray\" stroke=\"black\" cx=\"51\" cy=\"-59.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"51\" y=\"-54.58\" font-family=\"Times,serif\" font-size=\"14.00\">z</text>\n</g>\n<!-- y&#45;&gt;z -->\n<g id=\"edge2\" class=\"edge\">\n<title>y&#45;&gt;z</title>\n<path fill=\"none\" stroke=\"black\" d=\"M51,-159.31C51,-140.5 51,-110.28 51,-88.16\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54.5,-88.43 51,-78.43 47.5,-88.43 54.5,-88.43\"/>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f06e89061d0>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pyro.render_model(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7224816",
   "metadata": {},
   "source": [
    "## Semisupervised models\n",
    "\n",
    "Pyro allows semisupervised models by allowing different sets of `*args,**kwargs` to be passed to a model. You can render semisupervised models by passing a list of different tuples `model_args` and/or a list of different `model_kwargs` to denote the different ways you use a model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "da332f10",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(x, y=None):\n",
    "    with pyro.plate(\"N\", 2):\n",
    "        z = pyro.sample(\"z\", dist.Normal(0, 1))\n",
    "        y = pyro.sample(\"y\", dist.Normal(0, 1), obs=y)\n",
    "        pyro.sample(\"x\", dist.Normal(y + z, 1), obs=x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "20b69fb9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\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 8.1.0 (0)\n -->\n<!-- Pages: 1 -->\n<svg width=\"166pt\" height=\"173pt\"\n viewBox=\"0.00 0.00 166.00 173.25\" 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 169.25)\">\n<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-169.25 162,-169.25 162,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster_N</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"8,-8 8,-157.25 150,-157.25 150,-8 8,-8\"/>\n<text text-anchor=\"middle\" x=\"136\" y=\"-15.95\" font-family=\"Times,serif\" font-size=\"14.00\">N</text>\n</g>\n<!-- z -->\n<g id=\"node1\" class=\"node\">\n<title>z</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"115\" cy=\"-131.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"115\" y=\"-126.58\" font-family=\"Times,serif\" font-size=\"14.00\">z</text>\n</g>\n<!-- x -->\n<g id=\"node3\" class=\"node\">\n<title>x</title>\n<ellipse fill=\"gray\" stroke=\"black\" cx=\"79\" cy=\"-59.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"79\" y=\"-54.58\" font-family=\"Times,serif\" font-size=\"14.00\">x</text>\n</g>\n<!-- z&#45;&gt;x -->\n<g id=\"edge1\" class=\"edge\">\n<title>z&#45;&gt;x</title>\n<path fill=\"none\" stroke=\"black\" d=\"M106.65,-114.01C102.38,-105.71 97.08,-95.4 92.27,-86.04\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"94.99,-84.69 87.3,-77.4 88.77,-87.89 94.99,-84.69\"/>\n</g>\n<!-- y -->\n<g id=\"node2\" class=\"node\">\n<title>y</title>\n<defs>\n<linearGradient id=\"node2_l_0\" gradientUnits=\"userSpaceOnUse\" x1=\"16\" y1=\"-131.25\" x2=\"70\" y2=\"-131.25\" >\n<stop offset=\"0\" style=\"stop-color:gray;stop-opacity:1.;\"/>\n<stop offset=\"1\" style=\"stop-color:white;stop-opacity:1.;\"/>\n</linearGradient>\n</defs>\n<ellipse fill=\"url(#node2_l_0)\" stroke=\"black\" cx=\"43\" cy=\"-131.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"43\" y=\"-126.58\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n</g>\n<!-- y&#45;&gt;x -->\n<g id=\"edge2\" class=\"edge\">\n<title>y&#45;&gt;x</title>\n<path fill=\"none\" stroke=\"black\" d=\"M51.35,-114.01C55.62,-105.71 60.92,-95.4 65.73,-86.04\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"69.23,-87.89 70.7,-77.4 63.01,-84.69 69.23,-87.89\"/>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f06e819a950>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pyro.render_model(\n",
    "    model,\n",
    "    model_kwargs=[\n",
    "        {\"x\": torch.zeros(2)},\n",
    "        {\"x\": torch.zeros(2), \"y\": torch.zeros(2)},\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "837047a8",
   "metadata": {},
   "source": [
    "# Rendering deterministic variables\n",
    "\n",
    "Pyro allows deterministic variables to be defined using `pyro.deterministic`. These variables can be rendered by setting `render_deterministic=True` in `pyro.render_model` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d90dc8d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_deterministic(data):\n",
    "    sigma = pyro.param(\"sigma\", torch.tensor([1.]), constraint=constraints.positive)\n",
    "    mu = pyro.param(\"mu\", torch.tensor([0.]))\n",
    "    x = pyro.sample(\"x\", dist.Normal(mu, sigma))\n",
    "    log_y = pyro.sample(\"y\", dist.Normal(x, 1))\n",
    "    y = pyro.deterministic(\"y_deterministic\", log_y.exp())\n",
    "    with pyro.plate(\"N\", len(data)):\n",
    "        eps_z_loc = pyro.sample(\"eps_z_loc\", dist.Normal(0, 1))\n",
    "        z_loc = pyro.deterministic(\"z_loc\", eps_z_loc + x, event_dim=0)\n",
    "        pyro.sample(\"z\", dist.Normal(z_loc, y), obs=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6fcc43d8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\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 8.1.0 (0)\n -->\n<!-- Pages: 1 -->\n<svg width=\"563pt\" height=\"483pt\"\n viewBox=\"0.00 0.00 563.38 483.25\" 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 479.25)\">\n<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-479.25 559.38,-479.25 559.38,4 -4,4\"/>\n<g id=\"clust1\" class=\"cluster\">\n<title>cluster_N</title>\n<polygon fill=\"none\" stroke=\"black\" points=\"8,-8 8,-229.25 140,-229.25 140,-8 8,-8\"/>\n<text text-anchor=\"middle\" x=\"126\" y=\"-15.95\" font-family=\"Times,serif\" font-size=\"14.00\">N</text>\n</g>\n<!-- x -->\n<g id=\"node1\" class=\"node\">\n<title>x</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"201\" cy=\"-275.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"201\" y=\"-270.57\" font-family=\"Times,serif\" font-size=\"14.00\">x</text>\n</g>\n<!-- y -->\n<g id=\"node2\" class=\"node\">\n<title>y</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"231\" cy=\"-203.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"231\" y=\"-198.57\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n</g>\n<!-- x&#45;&gt;y -->\n<g id=\"edge1\" class=\"edge\">\n<title>x&#45;&gt;y</title>\n<path fill=\"none\" stroke=\"black\" d=\"M208.11,-257.66C211.55,-249.64 215.76,-239.82 219.62,-230.79\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"223.22,-232.29 223.94,-221.72 216.79,-229.53 223.22,-232.29\"/>\n</g>\n<!-- z_loc -->\n<g id=\"node7\" class=\"node\">\n<title>z_loc</title>\n<ellipse fill=\"gray\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"95\" cy=\"-131.25\" rx=\"35.42\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"95\" y=\"-126.58\" font-family=\"Times,serif\" font-size=\"14.00\">z_loc</text>\n</g>\n<!-- x&#45;&gt;z_loc -->\n<g id=\"edge3\" class=\"edge\">\n<title>x&#45;&gt;z_loc</title>\n<path fill=\"none\" stroke=\"black\" d=\"M190.94,-258.09C179.46,-239.93 159.82,-209.8 141,-185.25 133.47,-175.43 124.68,-165.06 116.81,-156.12\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"119.88,-154.19 110.62,-149.06 114.66,-158.85 119.88,-154.19\"/>\n</g>\n<!-- y_deterministic -->\n<g id=\"node3\" class=\"node\">\n<title>y_deterministic</title>\n<ellipse fill=\"gray\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"237\" cy=\"-131.25\" rx=\"88.56\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"237\" y=\"-126.58\" font-family=\"Times,serif\" font-size=\"14.00\">y_deterministic</text>\n</g>\n<!-- y&#45;&gt;y_deterministic -->\n<g id=\"edge2\" class=\"edge\">\n<title>y&#45;&gt;y_deterministic</title>\n<path fill=\"none\" stroke=\"black\" d=\"M232.48,-184.95C233.12,-177.49 233.89,-168.57 234.6,-160.22\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"238.17,-160.62 235.53,-150.35 231.19,-160.02 238.17,-160.62\"/>\n</g>\n<!-- z -->\n<g id=\"node8\" class=\"node\">\n<title>z</title>\n<ellipse fill=\"gray\" stroke=\"black\" cx=\"100\" cy=\"-59.25\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"100\" y=\"-54.58\" font-family=\"Times,serif\" font-size=\"14.00\">z</text>\n</g>\n<!-- y_deterministic&#45;&gt;z -->\n<g id=\"edge5\" class=\"edge\">\n<title>y_deterministic&#45;&gt;z</title>\n<path fill=\"none\" stroke=\"black\" d=\"M205.23,-114.01C182.74,-102.53 152.76,-87.21 130.56,-75.86\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"132.45,-72.39 121.96,-70.96 129.27,-78.62 132.45,-72.39\"/>\n</g>\n<!-- sigma -->\n<g id=\"node4\" class=\"node\">\n<title>sigma</title>\n<polygon fill=\"white\" stroke=\"none\" points=\"197.38,-410.88 154.62,-410.88 154.62,-393.62 197.38,-393.62 197.38,-410.88\"/>\n<text text-anchor=\"middle\" x=\"176\" y=\"-397.57\" font-family=\"Times,serif\" font-size=\"14.00\">sigma</text>\n</g>\n<!-- sigma&#45;&gt;x -->\n<g id=\"edge7\" class=\"edge\">\n<title>sigma&#45;&gt;x</title>\n<path fill=\"none\" stroke=\"black\" d=\"M177.49,-393.81C180.94,-376.54 189.64,-333.06 195.43,-304.1\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"199.04,-304.92 197.56,-294.43 192.17,-303.55 199.04,-304.92\"/>\n</g>\n<!-- mu -->\n<g id=\"node5\" class=\"node\">\n<title>mu</title>\n<polygon fill=\"white\" stroke=\"none\" points=\"238.25,-410.88 215.75,-410.88 215.75,-393.62 238.25,-393.62 238.25,-410.88\"/>\n<text text-anchor=\"middle\" x=\"227\" y=\"-397.57\" font-family=\"Times,serif\" font-size=\"14.00\">mu</text>\n</g>\n<!-- mu&#45;&gt;x -->\n<g id=\"edge8\" class=\"edge\">\n<title>mu&#45;&gt;x</title>\n<path fill=\"none\" stroke=\"black\" d=\"M225.45,-393.81C221.86,-376.54 212.82,-333.06 206.79,-304.1\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"210.04,-303.5 204.57,-294.43 203.18,-304.93 210.04,-303.5\"/>\n</g>\n<!-- eps_z_loc -->\n<g id=\"node6\" class=\"node\">\n<title>eps_z_loc</title>\n<ellipse fill=\"white\" stroke=\"black\" cx=\"74\" cy=\"-203.25\" rx=\"57.52\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"74\" y=\"-198.57\" font-family=\"Times,serif\" font-size=\"14.00\">eps_z_loc</text>\n</g>\n<!-- eps_z_loc&#45;&gt;z_loc -->\n<g id=\"edge4\" class=\"edge\">\n<title>eps_z_loc&#45;&gt;z_loc</title>\n<path fill=\"none\" stroke=\"black\" d=\"M79.19,-184.95C81.48,-177.32 84.23,-168.17 86.78,-159.65\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"90.35,-160.94 89.87,-150.35 83.64,-158.93 90.35,-160.94\"/>\n</g>\n<!-- z_loc&#45;&gt;z -->\n<g id=\"edge6\" class=\"edge\">\n<title>z_loc&#45;&gt;z</title>\n<path fill=\"none\" stroke=\"black\" d=\"M96.24,-112.95C96.77,-105.49 97.41,-96.57 98,-88.22\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"101.56,-88.58 98.78,-78.35 94.57,-88.08 101.56,-88.58\"/>\n</g>\n<!-- distribution_description_node -->\n<g id=\"node9\" class=\"node\">\n<title>distribution_description_node</title>\n<text text-anchor=\"start\" x=\"264.62\" y=\"-457.95\" font-family=\"Times,serif\" font-size=\"14.00\">x ~ Normal</text>\n<text text-anchor=\"start\" x=\"264.62\" y=\"-440.7\" font-family=\"Times,serif\" font-size=\"14.00\">y ~ Normal</text>\n<text text-anchor=\"start\" x=\"264.62\" y=\"-423.45\" font-family=\"Times,serif\" font-size=\"14.00\">y_deterministic ~ Deterministic</text>\n<text text-anchor=\"start\" x=\"264.62\" y=\"-406.2\" font-family=\"Times,serif\" font-size=\"14.00\">eps_z_loc ~ Normal</text>\n<text text-anchor=\"start\" x=\"264.62\" y=\"-388.95\" font-family=\"Times,serif\" font-size=\"14.00\">z_loc ~ Deterministic</text>\n<text text-anchor=\"start\" x=\"264.62\" y=\"-371.7\" font-family=\"Times,serif\" font-size=\"14.00\">z ~ Normal</text>\n<text text-anchor=\"start\" x=\"264.62\" y=\"-354.45\" font-family=\"Times,serif\" font-size=\"14.00\">sigma : GreaterThan(lower_bound=0.0)</text>\n<text text-anchor=\"start\" x=\"264.62\" y=\"-337.2\" font-family=\"Times,serif\" font-size=\"14.00\">mu : Real()</text>\n</g>\n</g>\n</svg>\n",
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f06e8a89f10>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = torch.ones(10)\n",
    "pyro.render_model(\n",
    "    model_deterministic,\n",
    "    model_args=(data,),\n",
    "    render_params=True,\n",
    "    render_distributions=True,\n",
    "    render_deterministic=True\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}