{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "861823e7",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "(Normalizing_Flows)=\n",
    "\n",
    "# Normalizing Flows in PyTensor\n",
    "\n",
    ":::{post} August 17, 2025 \n",
    ":tags: Graph rewrites \n",
    ":category: avanced, explanation \n",
    ":author: Jesse Grabowski, Ricardo Vieira\n",
    ":::\n",
    "\n",
    "\n",
    "A hip new algorithm for doing machine learning on distributions is called *normalizing flows*. \n",
    "\n",
    "The idea begins with the change of variables formula. We begin with a random variable $X \\sim N(0, 1)$. Suppose we know the PDF for $X$ (because we do!). What if we instead wanted the PDF for $ Y = \\exp(X)$. Do we know that? \n",
    "\n",
    "The answer is yes! We know it because of the *change of variable* formula, which states:\n",
    "\n",
    "$$\n",
    "g(y) = (f \\circ G^{-1})(y) \\left | \\frac{\\partial}{\\partial x} G^{-1}(y) \\right |\n",
    "$$\n",
    "\n",
    "Where:\n",
    "\n",
    "- $g(\\cdot)$ is the (unknown!) PDF of the variable $Y$\n",
    "- $f(\\cdot)$ if the (known!) PDF of the variable $X$\n",
    "- $G(\\cdot)$ is a function with nice properties.\n",
    "\n",
    "The \"nice properties\" require (in the most general case) that $G(x)$ is a $C^1$ diffeomorphism, which means that it is 1) continuous and differentiable almost everywhere; 2) it is bijective, and 3) its derivatives are also bijective. \n",
    "\n",
    "A simpler requirement is that $G(x)$ is continuous, bijective, and monotonic. That will get us 99% of the way there. Hey, $\\exp$ is continuous, bijective, and monotonic -- what a coincidence!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d61e7e32",
   "metadata": {},
   "source": [
    "## Prepare Notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "327c0f12",
   "metadata": {},
   "outputs": [],
   "source": [
    "from abc import ABC\n",
    "\n",
    "import arviz as az\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import preliz as pz\n",
    "import pymc as pm\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "import pytensor\n",
    "import pytensor.tensor as pt\n",
    "from pytensor.graph.basic import explicit_graph_inputs\n",
    "from pytensor.graph.replace import vectorize_graph\n",
    "from pytensor.graph.rewriting import rewrite_graph\n",
    "from pytensor.tensor.basic import TensorVariable\n",
    "\n",
    "\n",
    "plt.style.use(\"seaborn-v0_8\")\n",
    "\n",
    "%config InlineBackend.figure_format = \"retina\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d3acdd",
   "metadata": {},
   "source": [
    "## Simple Example: Change of Variables Formula"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f318527d",
   "metadata": {},
   "source": [
    "Let's visualize this result through a concrete example. We look at the transformation of a normal distribution to a lognormal distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "99ba1913",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/wj/wjy2vm8d7_j9v43bv29zcgl80000gq/T/ipykernel_79672/3343765610.py:11: RuntimeWarning: invalid value encountered in log\n",
      "  return np.where(x > 0, np.log(x), 0)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 73916 characters]",
      "text/plain": [
       "<Figure size 800x550 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 466,
       "width": 669
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "X = pz.Normal(0, 1)\n",
    "Y = pz.LogNormal(0, 1)\n",
    "\n",
    "x_grid = np.linspace(-3.1, 3.1, 100)\n",
    "\n",
    "f_values = X.pdf(x_grid)\n",
    "\n",
    "\n",
    "# This is just a convenient trick: setting the value to 0 if x <= 0\n",
    "def G_inv(x):\n",
    "    return np.where(x > 0, np.log(x), 0)\n",
    "\n",
    "\n",
    "def d_G_inv(x):\n",
    "    return np.where(x > 0, np.reciprocal(x), 0)\n",
    "\n",
    "\n",
    "g_values = (X.pdf(G_inv(x_grid))) * np.abs(d_G_inv(x_grid))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(x_grid, Y.pdf(x_grid), c=\"C0\", label=\"LogNormal PDF\")\n",
    "ax.plot(\n",
    "    x_grid, g_values, c=\"C1\", ls=\"--\", label=r\"$f(\\ln x) \\left | \\frac{1}{x} \\right |$\"\n",
    ")\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9ba9923",
   "metadata": {},
   "source": [
    "## Normalizing Flows Fundamentals\n",
    "\n",
    "In this section, we will explore the core concepts of normalizing flows. along this notebook, we will work one complete example: sampling from a log-normal distribution using normalizing flows.\n",
    "\n",
    "Before going into the actual example, let us understand the overall strategy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a4a265c",
   "metadata": {},
   "source": [
    "The change of variables procedure above, is basically all we need to understand normalizing flows. The logic goes like this:\n",
    "\n",
    "1. Suppose we have some data $\\mathcal{D} \\sim D(\\cdot)$, where $D$ is totally unknown.\n",
    "2. We would like to sample from $D$, which seems hard since we don't know it.\n",
    "3. Instead of sampling from $D$, let's instead sample $F(\\cdot)$, where we get to choose $F$ to be as simple as possible. Say it's just $N(0, 1)$\n",
    "4. Then, we'll use a change of variables to make draws from $F$ _look like_ draws from $D$. \n",
    "5. We just need a function $G(x)$ that makes *samples* $x \\sim N(0, 1)$ look like our data.\n",
    "6. $G$ needs to have the nice properties we want: monotonic, bijective, and invertible. And, we'd  like it to have some parameters to learn!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc400f29",
   "metadata": {},
   "source": [
    "Let's define a simple base class to represent normalizing flows in PyTensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "75a89714",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Flow(ABC):\n",
    "    __slots__ = (\"parameters\",)\n",
    "\n",
    "    def transform(self, x: TensorVariable) -> TensorVariable:\n",
    "        \"\"\"Apply function F to the distribution x, along with the jacobian correction\"\"\"\n",
    "        ...\n",
    "\n",
    "    def inverse_transform(self, x: TensorVariable) -> TensorVariable:\n",
    "        \"\"\"Apply inverse function F^{-1} to the distribtuion x\"\"\"\n",
    "        ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a6e041f",
   "metadata": {},
   "source": [
    "We can now specify two concrete examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b4770d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Exp(Flow):\n",
    "    \"\"\"Exponential transformation\"\"\"\n",
    "\n",
    "    def transform(self, x: TensorVariable) -> TensorVariable:\n",
    "        return pt.exp(x)\n",
    "\n",
    "    def inverse_transform(self, x: TensorVariable) -> TensorVariable:\n",
    "        return pt.log(x)\n",
    "\n",
    "\n",
    "class Affine(Flow):\n",
    "    \"\"\"Affine transformation\"\"\"\n",
    "\n",
    "    __slots__ = (\"loc\", \"scale\", \"parameters\")\n",
    "\n",
    "    def __init__(self, loc: float = 0.0, scale: float = 1.0) -> None:\n",
    "        self.loc = pt.as_tensor_variable(loc)\n",
    "        self.scale = pt.as_tensor_variable(scale)\n",
    "\n",
    "    def transform(self, x: TensorVariable) -> TensorVariable:\n",
    "        return self.loc + self.scale * x\n",
    "\n",
    "    def inverse_transform(self, x: TensorVariable) -> TensorVariable:\n",
    "        return (x - self.loc) / self.scale\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be5d7585",
   "metadata": {},
   "source": [
    "As the composition of two flows is itself a flow, we can define a class that composes them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "09e1c0dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compose(*fns):\n",
    "    def f(x):\n",
    "        for fn in fns:\n",
    "            x = fn(x)\n",
    "        return x\n",
    "\n",
    "    return f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d93829f7",
   "metadata": {},
   "source": [
    "## Understanding the Forward Pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66e7f3b2",
   "metadata": {},
   "source": [
    "Let's use these classes to generate a normalizing flow computational graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ace905a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exp [id A]\n",
      " └─ Add [id B]\n",
      "    ├─ loc [id C]\n",
      "    └─ Mul [id D]\n",
      "       ├─ scale [id E]\n",
      "       └─ normal_rv{\"(),()->()\"}.1 [id F]\n",
      "          ├─ rng [id G]\n",
      "          ├─ NoneConst{None} [id H]\n",
      "          ├─ 0 [id I]\n",
      "          └─ 1 [id J]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x1055e99c0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Random number generation\n",
    "rng = pt.random.shared_rng(seed=0, name=\"rng\")\n",
    "# Sample from a normal distribution\n",
    "next_rng, x = rng.normal(0, 1)\n",
    "\n",
    "# Define the parameters of the affine transformation\n",
    "loc, scale = pt.dscalars(\"loc\", \"scale\")\n",
    "# Define the transformations\n",
    "transforms = [Affine(loc=loc, scale=scale), Exp()]\n",
    "# Compose the transformations\n",
    "flow = compose(*[x.transform for x in transforms])\n",
    "# Apply the transformations to the random variable\n",
    "z = flow(x)\n",
    "\n",
    "# Print the resulting graph\n",
    "z.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "501932bf",
   "metadata": {},
   "source": [
    "We can now compute the log-density of the transformed variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "975b389d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Switch [id A] 'z_values_logprob'\n",
      " ├─ Isnan [id B]\n",
      " │  └─ Mul [id C]\n",
      " │     ├─ [-1.] [id D]\n",
      " │     └─ Log [id E]\n",
      " │        └─ z_values [id F]\n",
      " ├─ [-inf] [id G]\n",
      " └─ Add [id H]\n",
      "    ├─ Switch [id I]\n",
      "    │  ├─ Isnan [id J]\n",
      "    │  │  └─ Mul [id K]\n",
      "    │  │     ├─ [-1.] [id D]\n",
      "    │  │     └─ Log [id L]\n",
      "    │  │        └─ Abs [id M]\n",
      "    │  │           └─ Alloc [id N]\n",
      "    │  │              ├─ scale [id O]\n",
      "    │  │              └─ Squeeze{axis=0} [id P]\n",
      "    │  │                 └─ Shape [id Q]\n",
      "    │  │                    └─ Sub [id R]\n",
      "    │  │                       ├─ Log [id E]\n",
      "    │  │                       │  └─ ···\n",
      "    │  │                       └─ ExpandDims{axis=0} [id S]\n",
      "    │  │                          └─ loc [id T]\n",
      "    │  ├─ [-inf] [id G]\n",
      "    │  └─ Add [id U]\n",
      "    │     ├─ Check{sigma > 0} [id V]\n",
      "    │     │  ├─ Add [id W]\n",
      "    │     │  │  ├─ [-0.91893853] [id X]\n",
      "    │     │  │  └─ Mul [id Y]\n",
      "    │     │  │     ├─ [-0.5] [id Z]\n",
      "    │     │  │     └─ Pow [id BA]\n",
      "    │     │  │        ├─ True_div [id BB]\n",
      "    │     │  │        │  ├─ Sub [id R]\n",
      "    │     │  │        │  │  └─ ···\n",
      "    │     │  │        │  └─ ExpandDims{axis=0} [id BC]\n",
      "    │     │  │        │     └─ scale [id O]\n",
      "    │     │  │        └─ [2] [id BD]\n",
      "    │     │  └─ True [id BE]\n",
      "    │     └─ Mul [id K]\n",
      "    │        └─ ···\n",
      "    └─ Mul [id C]\n",
      "       └─ ···\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x1055e99c0>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_values = pt.dvector(\"z_values\")\n",
    "# The function `pm.logp` does the magic!\n",
    "z_logp = pm.logp(z, z_values, jacobian=True)\n",
    "# We do this rewrite to make the computation more stable.\n",
    "rewrite_graph(z_logp).dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "911e0bd9",
   "metadata": {},
   "source": [
    "From the log-density graph, we can extract the inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9ca6e924",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[z_values, scale, loc]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs = list(explicit_graph_inputs(z_logp))\n",
    "inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "49ef5ec0",
   "metadata": {},
   "outputs": [],
   "source": [
    "z_values, scale, loc = inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "448fa2fe",
   "metadata": {},
   "source": [
    "We proceed to compile the graph into a function so we can do actual computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a27b4303",
   "metadata": {},
   "outputs": [],
   "source": [
    "f_logp_pymc = pytensor.function([z_values, loc, scale], z_logp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cd883ba",
   "metadata": {},
   "source": [
    "Let's verify that the resulting log-density of the transformed variable is the same as the log-density of a `LogNormal` distribution, which is precisely what we expect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "24bb63d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 80028 characters]",
      "text/plain": [
       "<Figure size 800x550 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 466,
       "width": 684
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x_grid = np.linspace(-3, 10, 1000)\n",
    "plt.plot(\n",
    "    x_grid,\n",
    "    np.exp(f_logp_pymc(z_values=x_grid, loc=3, scale=1.0)),\n",
    "    c=\"C0\",\n",
    "    label=\"PyMC logp through the graph\",\n",
    ")\n",
    "plt.plot(\n",
    "    x_grid,\n",
    "    pz.LogNormal(mu=3, sigma=1).pdf(x_grid),\n",
    "    ls=\"--\",\n",
    "    c=\"C1\",\n",
    "    label=\"LogNormal PDF\",\n",
    ")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0a3d904",
   "metadata": {},
   "source": [
    "## Learning the `loc` and `scale` parameters\n",
    "\n",
    "Suppose now, we just have samples from some \"unknown\" lognormal distribution, we can learn parameters `loc` and `scale` by maximizing the log likelihood of the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "88dba68c",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = pz.LogNormal(mu=2.234, sigma=0.99354).rvs(5_000)\n",
    "\n",
    "objective = -pm.logp(z, z_values, jacobian=True).mean()\n",
    "\n",
    "# We can use pytensor for automatic differentiation\n",
    "grad_objective = pt.stack(pt.grad(objective, [loc, scale]))\n",
    "# We compile the graph to make actual computations\n",
    "f_obj_grad = pytensor.function([z_values, loc, scale], [objective, grad_objective])\n",
    "\n",
    "# Next, we compute the Hessian matrix\n",
    "# It would be nicer to just use the hessian function, but pytensor is fussy and doesn't like\n",
    "# scalar inputs\n",
    "hess = pt.stack(\n",
    "    pytensor.gradient.jacobian(\n",
    "        pt.stack(pt.grad(objective, [loc, scale])), [loc, scale]\n",
    "    ),\n",
    "    axis=0,\n",
    ")\n",
    "\n",
    "f_hess = pytensor.function([z_values, loc, scale], hess)\n",
    "\n",
    "f_sample = pm.compile([loc, scale], z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "446f2d5c",
   "metadata": {},
   "source": [
    "We can use the `scipy.optimize.minimize` function to find the maximum likelihood estimates of the `loc` and `scale` parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6704f27e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       " message: Optimization terminated successfully.\n",
       " success: True\n",
       "  status: 0\n",
       "     fun: 3.6165369785017196\n",
       "       x: [ 2.232e+00  9.664e-01]\n",
       "     nit: 10\n",
       "     jac: [ 1.643e-13 -4.547e-13]\n",
       "    nfev: 12\n",
       "    njev: 12\n",
       "    nhev: 10"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = minimize(\n",
    "    lambda x, *args: f_obj_grad(*args, *x),\n",
    "    jac=True,\n",
    "    hess=lambda x, *args: f_hess(*args, *x),\n",
    "    method=\"Newton-CG\",\n",
    "    x0=[0.8, 0.8],\n",
    "    args=(samples,),\n",
    "    tol=1e-12,\n",
    ")\n",
    "res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79aa4f16",
   "metadata": {},
   "source": [
    "We see that we have recovered the parameters of the lognormal distribution. Let's compare the samples from the learned distribution with the original samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "feb5b9db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 138856 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 557,
       "width": 860
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "az.plot_kde(\n",
    "    np.r_[*[f_sample(*res.x) for _ in range(5_000)]],\n",
    "    plot_kwargs={\"c\": \"C0\", \"label\": \"Samples from Learned Distribution\"},\n",
    "    ax=ax,\n",
    ")\n",
    "az.plot_kde(\n",
    "    samples,\n",
    "    plot_kwargs={\"c\": \"C1\", \"label\": \"Samples from Original Distribution\"},\n",
    "    ax=ax,\n",
    ")\n",
    "ax.legend(fontsize=12)\n",
    "ax.set_xlabel(\"x\", fontsize=14)\n",
    "ax.set_ylabel(\"Density\", fontsize=14)\n",
    "ax.tick_params(labelsize=12)\n",
    "ax.set_title(\n",
    "    \"Samples from LogNormal(mu=2.234, sigma=0.99354)\", fontsize=18, fontweight=\"bold\"\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f9a7a50",
   "metadata": {},
   "source": [
    "These distribution are essentially the same."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e8ce3cb",
   "metadata": {},
   "source": [
    "## Reproducing what PyMC does\n",
    "\n",
    "Most of the heavy lifting here was done by PyMC. It already knows how to handle the change-of-variables formula for a wide variety of functions, including $\\exp$ and affine shifts. That's why things just sort of magically worked up until now. \n",
    "\n",
    "Recall the formula for change of variables:\n",
    "\n",
    "$$\n",
    "g(x) = (f \\circ G^{-1})(x) \\left | \\frac{\\partial}{\\partial x} G^{-1}(x) \\right |\n",
    "$$\n",
    "\n",
    "Right now in the graph of $z$, we just have draws from $x$, passed through the function $G$. Or:\n",
    "\n",
    "$$\n",
    "z = G(x) \n",
    "$$\n",
    "\n",
    "To get to where we want to go, we need to:\n",
    "\n",
    "1. Iterate backwards over all the flows, applying the inverse function to the incoming samples\n",
    "2. At the same time, we also need to compute $\\left | \\frac{\\partial}{\\partial x} G^{-1}(x) \\right | $ and accumulate their products\n",
    "3. We need to get the logp of the original random variable, $f(x)$, and plug in the result of (1) as the value\n",
    "4. Finally, add the log of the product of the determinant corrections. \n",
    "\n",
    "Since in the end we need the log of the determinant, it will be easier to modify step (2) to accumulate the sum of the log determinant of the transformations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "207209eb",
   "metadata": {},
   "source": [
    "### In detail\n",
    "\n",
    "We have two transformations, $H(x) = \\exp(x)$ and $J(x) = a + bx$. Define $G$ as their composition, $G \\equiv (H \\circ J)(x) = \\exp(a + bx)$\n",
    "\n",
    "For inverses, we have $H^{-1}(x) = \\ln(x)$ and $J^{-1}(x) = \\frac{x - a}{b}$\n",
    "\n",
    "So, the inverse of their composition is $G^{-1} \\equiv (J^{-1} \\circ H^{-1}) = J^{-1}(H^{-1}(x)) = J^{-1}(\\ln(x)) = \\frac{\\ln(x) - a}{b}$\n",
    "\n",
    "For the correction term, we need the determinant of the jacobian. Since $G$ is a scalar function, this is just the absolute value of the gradient:\n",
    "\n",
    "$$\\left | \\frac{\\partial}{\\partial x}G^{-1} \\right | = \\left | \\frac{\\partial}{\\partial x} \\frac{\\ln(x) - a}{b} \\right | = \\left | \\frac{1}{b} \\cdot \\frac{1}{x} \\right | $$\n",
    "\n",
    "We we will compute $g(z) = f(\\frac{\\ln(z) - a}{b}) \\left | \\frac{1}{b} \\cdot \\frac{1}{z} \\right |$, where $f(x) = \\frac{1}{\\sqrt{2\\pi}}\\exp(-\\frac{x^2}{2})$ is a standard normal PDF.\n",
    "\n",
    "Of course we'll work with the logp, so we will actually have:\n",
    "\n",
    "$$\\ln g(z) = -\\frac{1}{2}\\ln{\\sqrt{2\\pi}} - \\frac{1}{2}\\left (\\frac{\\ln z - a}{b} \\right )^2 - \\ln{\\left | b \\right |} - \\ln{\\left | z \\right |}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05bc8bd2",
   "metadata": {},
   "source": [
    "### Solution by hand\n",
    "\n",
    "We now implement this analytic procedure in PyTensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ea6f42c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def standard_normal_logp(x):\n",
    "    x = pt.as_tensor_variable(x)\n",
    "\n",
    "    kernel = -(x**2) / 2\n",
    "    normalizing_constant = -0.5 * pt.log(2 * np.pi)\n",
    "    return normalizing_constant + kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f1736b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the parameters of the first transform (affine)\n",
    "loc, scale = transforms[0].loc, transforms[0].scale\n",
    "z_star = (pt.log(z_values) - loc) / scale\n",
    "analytic = (\n",
    "    standard_normal_logp(z_star) - pt.log(pt.abs(scale)) - pt.log(pt.abs(z_values))\n",
    ")\n",
    "# Compile the function\n",
    "f_analytic = pytensor.function([z_values, loc, scale], analytic)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "786f3be8",
   "metadata": {},
   "source": [
    "We can pass specific value now and evaluate the analytic function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "85d583bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-4.06681489, -3.42976537, -3.42097081, ..., -3.73252572,\n",
       "       -4.26441774, -3.43616604], shape=(5000,))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_analytic(samples, loc=3, scale=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcd081d3",
   "metadata": {},
   "source": [
    "We can verify these values are exactly what we are expecting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7eb42be1",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_almost_equal(\n",
    "    actual=f_analytic(samples, loc=3, scale=1),\n",
    "    desired=pz.LogNormal(mu=3, sigma=1).logpdf(samples),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10633413",
   "metadata": {},
   "source": [
    "### Solution using autodiff\n",
    "\n",
    "Now we do it again using autodiff. We need the absolute value of the gradient through the whole composite transformation. We do it on a dummy scalar $x$, then vectorize back up to `z_values`, since that's the behavior we want.\n",
    "\n",
    "We could also do `pt.grad(f_inverse(z_values).sum(), z_values)`, and rely on the fact that the sum function introduces no cross-terms in the gradient expression. This would be the same, but would require you to grok this math fact. We also like that the vectorize way shows off `vectorize_graph`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1e6e9c07",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = pt.dscalar(\"x\")\n",
    "# Get the inverse of the flows\n",
    "inverse_flows = [f.inverse_transform for f in reversed(transforms)]\n",
    "\n",
    "f_inverse = compose(*inverse_flows)\n",
    "# Use autodiff to compute the log jacobian determinant\n",
    "log_jac_det = pt.log(pt.abs(pt.grad(f_inverse(x), x)))\n",
    "# Vectorize the graph\n",
    "log_jac_det = vectorize_graph(log_jac_det, {x: z_values})\n",
    "\n",
    "# Compute the logp of the inverse flow\n",
    "inverse_flow = standard_normal_logp(f_inverse(z_values)) + log_jac_det\n",
    "# Compile the function\n",
    "f_logp_flow = pytensor.function(list(explicit_graph_inputs(inverse_flow)), inverse_flow)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46834a6f",
   "metadata": {},
   "source": [
    "As above let's verify that the results are consistent and correct:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "8323642a",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_almost_equal(\n",
    "    actual=f_logp_flow(samples, loc=3, scale=1),\n",
    "    desired=f_logp_pymc(samples, loc=3, scale=1),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abeea690",
   "metadata": {},
   "source": [
    "### Re-do optimization using our new implementation\n",
    "\n",
    "Finally, we can re-do the optimization using our new implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "4f6f537b",
   "metadata": {},
   "outputs": [],
   "source": [
    "objective = -inverse_flow.mean()\n",
    "\n",
    "grad_objective = pt.stack(pt.grad(objective, [loc, scale]))\n",
    "f_obj_grad = pytensor.function([z_values, loc, scale], [objective, grad_objective])\n",
    "\n",
    "hess = pt.stack(\n",
    "    pytensor.gradient.jacobian(\n",
    "        pt.stack(pt.grad(objective, [loc, scale])), [loc, scale]\n",
    "    ),\n",
    "    axis=0,\n",
    ")\n",
    "f_hess = pytensor.function([z_values, loc, scale], hess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "6f1f5acb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       " message: Optimization terminated successfully.\n",
       " success: True\n",
       "  status: 0\n",
       "     fun: 3.616536978501722\n",
       "       x: [ 2.232e+00  9.664e-01]\n",
       "     nit: 10\n",
       "     jac: [ 1.643e-13 -4.357e-13]\n",
       "    nfev: 12\n",
       "    njev: 12\n",
       "    nhev: 10"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = minimize(\n",
    "    lambda x, *args: f_obj_grad(*args, *x),\n",
    "    jac=True,\n",
    "    hess=lambda x, *args: f_hess(*args, *x),\n",
    "    method=\"Newton-CG\",\n",
    "    x0=[0.8, 0.8],\n",
    "    args=(samples,),\n",
    "    tol=1e-12,\n",
    ")\n",
    "res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "854383a5",
   "metadata": {},
   "source": [
    "We get the expected results 🚀!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea866a18",
   "metadata": {},
   "source": [
    "## Authors\n",
    "\n",
    "- Authored by Jesse Grabowski and Ricardo Vieira in August 2025"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2a10b15",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    ":::{bibliography} :filter: docname in docnames"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb49ec8",
   "metadata": {},
   "source": [
    "## Watermark "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "29f26df2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last updated: Sun Aug 17 2025\n",
      "\n",
      "Python implementation: CPython\n",
      "Python version       : 3.12.11\n",
      "IPython version      : 9.4.0\n",
      "\n",
      "pytensor: 2.31.7\n",
      "\n",
      "numpy     : 2.2.6\n",
      "pytensor  : 2.31.7\n",
      "arviz     : 0.22.0\n",
      "preliz    : 0.20.0\n",
      "pymc      : 5.25.1\n",
      "scipy     : 1.16.1\n",
      "matplotlib: 3.10.5\n",
      "\n",
      "Watermark: 2.5.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark -n -u -v -iv -w -p pytensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3e24733",
   "metadata": {},
   "source": [
    ":::{include} ../page_footer.md \n",
    ":::"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
