{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "(dimensionality)=\n",
    "# Distribution Dimensionality\n",
    "PyMC provides a number of ways to specify the dimensionality of its distributions. This document provides an overview, and offers some user tips.\n",
    "\n",
    "## Glossary\n",
    "In this document we'll be using the term dimensionality to refer to the idea of dimensions. Each of the terms below has a specific\n",
    "semantic and computational definition in PyMC.  While we share them here they will make much more sense when viewed in the examples below.\n",
    "\n",
    "+ *Support dimensions* → The core dimensionality of a distribution\n",
    "+ *Batch dimensions* → Extra dimensions beyond the support dimensionality of a distribution\n",
    "+ *Implicit dimensions* → Dimensions that follow from the values or shapes of the distribution parameters\n",
    "+ *Explicit dimensions* → Dimensions that are explicitly defined by one of the following arguments:\n",
    "    + *Shape* → Number of draws from a distribution\n",
    "    + *Dims* → An array of dimension names\n",
    "+ *Coords* → A dictionary mapping dimension names to coordinate values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "import pymc as pm\n",
    "import numpy as np\n",
    "import pytensor.tensor as pt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Univariate distribution example\n",
    "We can start with the simplest case, a single Normal distribution. We use `.dist` to specify one outside of a PyMC Model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "normal_dist = pm.Normal.dist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can then use the {func}`~pymc.draw` function to take a random draw from that distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Just patching the draw function for reproducibility\n",
    "rng = np.random.default_rng(seed=sum(map(ord, \"dimensionality\")))\n",
    "draw = partial(pm.draw, random_seed=rng)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array(0.80189558), 0)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normal_draw = draw(normal_dist)\n",
    "normal_draw, normal_draw.ndim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In this case we end up with a single scalar value. This means that a Normal distribution has a scalar support dimensionality, as the smallest random draw you can take is a scalar which has a dimension of zero. The support dimensionality of every distribution is hard-coded as a property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normal_dist.owner.op.ndim_supp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Explicit batch dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "If one needs more than a single draw, a natural tendency would be to create multiple copies of the same variable and stack them together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.83636296, -0.33327414,  0.9434115 ])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normal_dists = pm.math.stack([pm.Normal.dist() for _ in range(3)])\n",
    "draw(normal_dists)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "More simply, one can create a *batch* of independent draws from the same distribution family by using the shape argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.98810294, -0.07003785, -0.37962748])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normal_dists = pm.Normal.dist(shape=(3,))\n",
    "draw(normal_dists)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 7.99932116e-04, -1.94407945e+00,  3.90451962e-01],\n",
       "       [ 1.10657367e+00,  6.49042149e-01, -1.09450185e+00],\n",
       "       [-2.96226305e-01,  1.41884595e+00, -1.31589441e+00],\n",
       "       [ 1.53109449e+00, -7.73771737e-01,  2.37418367e+00]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normal_dists = pm.Normal.dist(shape=(4, 3))\n",
    "draw(normal_dists)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Not only is this more succinct, but it produces much more efficient vectorized code. We rarely use the stack approach in PyMC, unless we need to combine draws from distinct distribution families."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Implicit batch dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "It is also possible to create a batch of draws by passing parameters with higher dimensions, without having to specify shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.81833093, -0.2891973 ,  1.2399946 ])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normal_dists = pm.Normal.dist(mu=np.array([0, 0, 0]), sigma=np.array([1, 1, 1]))\n",
    "draw(normal_dists)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "This is equivalent to the previous example with explicit shape, and we could have passed it explicitly here. Because we did not, we refer to these batch dimensions as being *implicit*.\n",
    "\n",
    "Where this becomes very useful is when we want the parameters to vary across batch dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  0.99989975,  10.00009874, 100.00004215])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw(pm.Normal.dist(mu=[1, 10, 100], sigma=0.0001))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "When the parameters don't have the same shapes, they are broadcasted, in a similar way to how NumPy works. In this case `sigma` was broadcasted to match the shape of `mu`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([  1,  10, 100]), array([0.0001, 0.0001, 0.0001])]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.broadcast_arrays([1, 10, 100], 0.0001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's important to understand how NumPy {ref}`broadcasting <numpy:basics.broadcasting>` works. When you do something that is not valid, you will easily encounter this sort of errors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (3,) and arg 1 with shape (2,).\n",
      "Apply node that caused the error: normal_rv{0, (0, 0), floatX, True}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F6427F8CAC0>), [], 11, [  1  10 100], [0.1 0.1])\n",
      "Toposort index: 0\n",
      "Inputs types: [RandomGeneratorType, TensorType(int64, shape=(0,)), TensorType(int64, shape=()), TensorType(int64, shape=(3,)), TensorType(float64, shape=(2,))]\n",
      "Inputs shapes: ['No shapes', (0,), (), (3,), (2,)]\n",
      "Inputs strides: ['No strides', (0,), (), (8,), (8,)]\n",
      "Inputs values: [Generator(PCG64) at 0x7F6427F8CAC0, array([], dtype=int64), array(11), array([  1,  10, 100]), array([0.1, 0.1])]\n",
      "Outputs clients: [['output'], ['output']]\n",
      "\n",
      "HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.\n",
      "HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    # shapes of (3,) and (2,) can't be broadcasted together\n",
    "    draw(pm.Normal.dist(mu=[1, 10, 100], sigma=[0.1, 0.1]))\n",
    "except ValueError as error:\n",
    "    print(error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combining implicit and explicit batch dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can combine explicit shape dimensions with implicit batch dimensions. As mentioned above, they can provide the same information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.49526775, -0.94608062,  1.66397913])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normal_dists = pm.Normal.dist(mu=np.array([0, 1, 2]), sigma=1, shape=(3,))\n",
    "draw(normal_dists)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But shape can also be used to extend beyond any implicit batch dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 2.22626513,  2.12938134,  0.49074886],\n",
       "       [ 0.08312601,  1.05049093,  1.91718083],\n",
       "       [-0.68191815,  1.43771096,  1.76780399],\n",
       "       [-0.59883241,  0.26954893,  2.74319335]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normal_dists = pm.Normal.dist(mu=np.array([0, 1, 2]), sigma=1, shape=(4, 3))\n",
    "draw(normal_dists)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that, due to broadcasting rules, explicit batch dimensions must always \"go on the left\" of any implicit dimensions. So in the previous example `shape=(4, 3)` is valid, but `shape=(3, 4)` is not, because the `mu` parameter can be broadcasted to the first shape but not to the second."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (3, 4) and arg 1 with shape (3,).\n",
      "Apply node that caused the error: normal_rv{0, (0, 0), floatX, True}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F64280725E0>), [3 4], 11, [0 1 2], 1.0)\n",
      "Toposort index: 0\n",
      "Inputs types: [RandomGeneratorType, TensorType(int64, shape=(2,)), TensorType(int64, shape=()), TensorType(int64, shape=(3,)), TensorType(float64, shape=())]\n",
      "Inputs shapes: ['No shapes', (2,), (), (3,), ()]\n",
      "Inputs strides: ['No strides', (8,), (), (8,), ()]\n",
      "Inputs values: [Generator(PCG64) at 0x7F64280725E0, array([3, 4]), array(11), array([0, 1, 2]), array(1.)]\n",
      "Outputs clients: [['output'], ['output']]\n",
      "\n",
      "HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.\n",
      "HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    draw(pm.Normal.dist(mu=np.array([0, 1, 2]), sigma=1, shape=(3, 4)))\n",
    "except ValueError as error:\n",
    "    print(error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you needed the Normal variables to have `shape=(3, 4)`, you can transpose it after defining it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.73397401, -0.18717845, -0.78548049,  1.64478883],\n",
       "       [ 3.54543846,  1.22954216,  2.13674063,  1.94194106],\n",
       "       [ 0.85294471,  3.52041332,  2.94428975,  3.25944187]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transposed_normals = pm.Normal.dist(mu=np.array([0, 1, 2]), sigma=1, shape=(4, 3)).T\n",
    "draw(transposed_normals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip} It's important not to confuse dimensions set in the definition of a distribution versus those set in downstream manipulations like transposition, indexing or broadcasting. When sampling with PyMC (be it via forward sampling or MCMC), the random draws will always emanate from the distribution shape. Notice how in the following example, a different number of \"random\" draws were actually taken, despite the two variables having the same final shape.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_normals = pm.Normal.dist(shape=(3,))\n",
    "broadcasted_normal = pt.broadcast_to(pm.Normal.dist(), (3,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([-0.45755879,  1.59975702,  0.20546749]),\n",
       " array([0.29866199, 0.29866199, 0.29866199]))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw(vector_normals), draw(broadcasted_normal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Multivariate distribution example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Some distributions by definition return more than one value when evaluated. This may be a vector of values or a matrix or an arbitrary multidimensional tensor. An example is the Multivariate Normal, which always returns a vector (an array with one dimension)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0.55390975, 2.17440418, 1.83014764]), 1)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mvnormal_dist = pm.MvNormal.dist(mu=np.ones(3), cov=np.eye(3))\n",
    "mvnormal_draw = draw(mvnormal_dist)\n",
    "mvnormal_draw, mvnormal_draw.ndim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As with any distribution, the support dimensionality is specified as a fixed property"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mvnormal_dist.owner.op.ndim_supp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Even if you specify a MvNormal with a single dimension, you get back a vector!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([-0.68893796]), 1)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smallest_mvnormal_dist = pm.MvNormal.dist(mu=[1], cov=[[1]])\n",
    "smallest_mvnormal_draw = draw(smallest_mvnormal_dist)\n",
    "smallest_mvnormal_draw, smallest_mvnormal_draw.ndim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Implicit support dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In the MvNormal examples we just saw, the support dimension was actually implicit. Nowhere did we specify we wanted a vector of 3 or 1 draws. This was inferred from the shape of `mu` and `cov`. As such, we refer to it as being an *implicit support dimension*. We could be a bit more explicit by using shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.57262853, 0.34230354, 1.96818163])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "explicit_mvnormal = pm.MvNormal.dist(mu=np.ones(3), cov=np.eye(3), shape=(3,))\n",
    "draw(explicit_mvnormal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{warning} However, note that at the time of writing shape is simply ignored for support dimensions. It serves merely as a \"type-hint\" for labeling the expected dimensions.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.0623799 , 0.84622693, 0.34046237])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ignored_shape_mvnormal = pm.MvNormal.dist(mu=np.ones(3), cov=np.eye(3), shape=(4,))\n",
    "draw(ignored_shape_mvnormal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Explicit batch dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As with univariate distributions, we can add explicit batched dimensions. We will use another vector distribution to illustrate this: the Multinomial. The following snippet defines a matrix of five independent Multinomial distributions, each of which is a vector of size 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[2, 0, 3],\n",
       "       [1, 1, 3],\n",
       "       [0, 2, 3],\n",
       "       [0, 1, 4],\n",
       "       [1, 0, 4]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw(pm.Multinomial.dist(n=5, p=[0.1, 0.3, 0.6], shape=(5, 3)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    ":::{warning} Again, note that shape has no effect on the support dimensionality\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 4],\n",
       "       [0, 0, 5],\n",
       "       [3, 1, 1],\n",
       "       [0, 1, 4],\n",
       "       [0, 2, 3]])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw(pm.Multinomial.dist(n=5, p=[0.1, 0.3, 0.6], shape=(5, 4)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the same reason, you must always define explicit batched dimensions \"to the left\" of the support dimension. The following will not behave as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[2, 0, 3],\n",
       "       [1, 3, 1],\n",
       "       [1, 1, 3]])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw(pm.Multinomial.dist(n=5, p=[0.1, 0.3, 0.6], shape=(3, 5)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you needed the Multinomial variables to have `shape=(3, 5)` you can transpose it after defining it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 0, 0, 0],\n",
       "       [2, 2, 1, 0, 3],\n",
       "       [3, 3, 4, 5, 2]])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transposed_multinomials = pm.Multinomial.dist(n=5, p=[0.1, 0.3, 0.6], shape=(5, 3)).T\n",
    "draw(transposed_multinomials)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Implicit batch dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As with univariate distributions, we can use different parameters for each batched dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 2, 2],\n",
       "       [0, 3, 7]])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "multinomial_dist = pm.Multinomial.dist(n=[5, 10], p=[0.1, 0.3, 0.6])\n",
    "draw(multinomial_dist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Which is equivalent to the more verbose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[2, 2, 1],\n",
       "       [0, 3, 7]])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw(pm.Multinomial.dist(n=[5, 10], p=[[0.1, 0.3, 0.6], [0.1, 0.3, 0.6]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "If you are familiar with NumPy broadcasting rules you may be curious of how does PyMC make this work. Naive broadcasting wouldn't work here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (2,) and arg 1 with shape (3,).\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    np.broadcast_arrays([5, 10], [0.1, 0.3, 0.6])\n",
    "except ValueError as exc:\n",
    "    print(exc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To understand what is going on, we need to introduce the concept of parameter core dimensions. The core dimensions of a distribution's parameter are the minimum number of dimensions the parameters need to have in order to define a distribution. In the Multinomial distribution, `n` must at least be an scalar integer, but `p` must be at least a vector that represents the probability of having an outcome on each category. So, for the Multinomial distribution, `n` has 0 core dimensions, and `p` has 1 core dimension. \n",
    "\n",
    "So if we have a vector of two `n`, we should actually broadcast the vector of `p` into a matrix with two such vectors, and pair each `n` with each broadcasted row of `p`. This works exactly like `np.vectorize`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">> 5 [0.1 0.3 0.6]\n",
      ">> 10 [0.1 0.3 0.6]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[1, 0, 4],\n",
       "       [1, 2, 7]])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def core_multinomial(n, p):\n",
    "    print(\">>\", n, p)\n",
    "    return draw(pm.Multinomial.dist(n, p))\n",
    "\n",
    "\n",
    "vectorized_multinomial = np.vectorize(core_multinomial, signature=\"(),(p)->(p)\")\n",
    "vectorized_multinomial([5, 10], [0.1, 0.3, 0.6])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The core dimensionality of each distribution parameter is also hard-coded as a property of each distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 1)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "multinomial_dist.owner.op.ndims_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implicit batch dimensions must still respect broadcasting rules. The following example is not valid because `n` has batched dimensions of `shape=(2,)` and `p` has batched dimensions of `shape=(3,)` which cannot be broadcasted together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "operands could not be broadcast together with remapped shapes [original->remapped]: (2,)  and requested shape (3,)\n",
      "Apply node that caused the error: multinomial_rv{1, (0, 1), int64, True}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F6425B8B060>), [], 4, [ 5 10], [[0.1 0.3 ... 0.3 0.6]])\n",
      "Toposort index: 0\n",
      "Inputs types: [RandomGeneratorType, TensorType(int64, shape=(0,)), TensorType(int64, shape=()), TensorType(int64, shape=(2,)), TensorType(float64, shape=(3, 3))]\n",
      "Inputs shapes: ['No shapes', (0,), (), (2,), (3, 3)]\n",
      "Inputs strides: ['No strides', (0,), (), (8,), (24, 8)]\n",
      "Inputs values: [Generator(PCG64) at 0x7F6425B8B060, array([], dtype=int64), array(4), array([ 5, 10]), 'not shown']\n",
      "Outputs clients: [['output'], ['output']]\n",
      "\n",
      "HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.\n",
      "HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    draw(pm.Multinomial.dist(n=[5, 10], p=[[0.1, 0.3, 0.6], [0.1, 0.3, 0.6], [0.1, 0.3, 0.6]]))\n",
    "except ValueError as error:\n",
    "    print(error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Combining implicit and explicit batch dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "You can and should combine implicit dimensions from multidimensional parameters with explicit shape information, which is easier to reason about."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 4],\n",
       "       [4, 1, 5]])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw(pm.Multinomial.dist(n=[5, 10], p=[0.1, 0.3, 0.6], shape=(2, 3)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Explicit batch dimensions can still extend beyond any implicit batch dimensions. Again, due to how broadcasting works, explicit batch dimensions must always \"go on the left\". The following case is invalid, because `n` has batched dimensions of `shape=(2,)`, which cannot be broadcasted to the explicit batch dimensions of `shape=(2, 4)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "operands could not be broadcast together with remapped shapes [original->remapped]: (2,)  and requested shape (2,4)\n",
      "Apply node that caused the error: multinomial_rv{1, (0, 1), int64, True}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F6425AC8120>), [2 4], 4, [ 5 10], [0.1 0.3 0.6])\n",
      "Toposort index: 0\n",
      "Inputs types: [RandomGeneratorType, TensorType(int64, shape=(2,)), TensorType(int64, shape=()), TensorType(int64, shape=(2,)), TensorType(float64, shape=(3,))]\n",
      "Inputs shapes: ['No shapes', (2,), (), (2,), (3,)]\n",
      "Inputs strides: ['No strides', (8,), (), (8,), (8,)]\n",
      "Inputs values: [Generator(PCG64) at 0x7F6425AC8120, array([2, 4]), array(4), array([ 5, 10]), array([0.1, 0.3, 0.6])]\n",
      "Outputs clients: [['output'], ['output']]\n",
      "\n",
      "HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.\n",
      "HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    draw(pm.Multinomial.dist(n=[5, 10], p=[0.1, 0.3, 0.6], shape=(2, 4, 3)))\n",
    "except ValueError as error:\n",
    "    print(error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Inspecting dimensionality with a model graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More often than not distributions are used inside a PyMC model, and as such there are tools that facilitate reasoning about distributions shapes in that context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          x: shape=(3,)\n",
      "sigma_log__: shape=()\n",
      "      sigma: shape=()\n",
      "          y: shape=(3,)\n"
     ]
    }
   ],
   "source": [
    "with pm.Model() as pmodel:\n",
    "    mu = pm.Normal(\"x\", mu=0, shape=(3))\n",
    "    sigma = pm.HalfNormal(\"sigma\")\n",
    "    y = pm.Normal(\"y\", mu=mu, sigma=sigma)\n",
    "\n",
    "for rv, shape in pmodel.eval_rv_shapes().items():\n",
    "    print(f\"{rv:>11}: shape={shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "An even more powerful tool to understand and debug dimensionality in PyMC is the {func}`~pymc.model_to_graphviz` function. Rather than inspecting array outputs we can instead read the Graphviz output to understand the dimensionality of the variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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 7.1.0 (20230122.1345)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"282pt\" height=\"249pt\"\n",
       " viewBox=\"0.00 0.00 281.71 248.91\" 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 244.91)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-244.91 277.71,-244.91 277.71,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>cluster3</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M20,-8C20,-8 110,-8 110,-8 116,-8 122,-14 122,-20 122,-20 122,-220.91 122,-220.91 122,-226.91 116,-232.91 110,-232.91 110,-232.91 20,-232.91 20,-232.91 14,-232.91 8,-226.91 8,-220.91 8,-220.91 8,-20 8,-20 8,-14 14,-8 20,-8\"/>\n",
       "<text text-anchor=\"middle\" x=\"109\" y=\"-15.8\" font-family=\"Times,serif\" font-size=\"14.00\">3</text>\n",
       "</g>\n",
       "<!-- y -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>y</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"65\" cy=\"-76.48\" rx=\"49.49\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
       "</g>\n",
       "<!-- x -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>x</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"65\" cy=\"-187.43\" rx=\"49.49\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-198.73\" font-family=\"Times,serif\" font-size=\"14.00\">x</text>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-183.73\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-168.73\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
       "</g>\n",
       "<!-- x&#45;&gt;y -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>x&#45;&gt;y</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M65,-149.54C65,-141.91 65,-133.76 65,-125.79\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"68.5,-125.84 65,-115.84 61.5,-125.84 68.5,-125.84\"/>\n",
       "</g>\n",
       "<!-- sigma -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>sigma</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"203\" cy=\"-187.43\" rx=\"70.92\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"203\" y=\"-198.73\" font-family=\"Times,serif\" font-size=\"14.00\">sigma</text>\n",
       "<text text-anchor=\"middle\" x=\"203\" y=\"-183.73\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"203\" y=\"-168.73\" font-family=\"Times,serif\" font-size=\"14.00\">HalfNormal</text>\n",
       "</g>\n",
       "<!-- sigma&#45;&gt;y -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>sigma&#45;&gt;y</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M164.19,-155.79C146.52,-141.84 125.55,-125.28 107.45,-110.99\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"110.03,-108.57 100.02,-105.12 105.7,-114.07 110.03,-108.57\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f64258f4a00>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pm.model_to_graphviz(pmodel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In the example above the number on the bottom left of each box (or plate) indicates the dimensionality of the distributions within. If a distribution is outside of any box with a number, it has a scalar shape.\n",
    "\n",
    "Let's use this tool to review implicit and explicit dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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 7.1.0 (20230122.1345)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"618pt\" height=\"138pt\"\n",
       " viewBox=\"0.00 0.00 617.92 137.95\" 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 133.95)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-133.95 613.92,-133.95 613.92,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>cluster3</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M205.92,-8C205.92,-8 383.92,-8 383.92,-8 389.92,-8 395.92,-14 395.92,-20 395.92,-20 395.92,-109.95 395.92,-109.95 395.92,-115.95 389.92,-121.95 383.92,-121.95 383.92,-121.95 205.92,-121.95 205.92,-121.95 199.92,-121.95 193.92,-115.95 193.92,-109.95 193.92,-109.95 193.92,-20 193.92,-20 193.92,-14 199.92,-8 205.92,-8\"/>\n",
       "<text text-anchor=\"middle\" x=\"382.92\" y=\"-15.8\" font-family=\"Times,serif\" font-size=\"14.00\">3</text>\n",
       "</g>\n",
       "<g id=\"clust2\" class=\"cluster\">\n",
       "<title>cluster4</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M415.92,-8C415.92,-8 589.92,-8 589.92,-8 595.92,-8 601.92,-14 601.92,-20 601.92,-20 601.92,-109.95 601.92,-109.95 601.92,-115.95 595.92,-121.95 589.92,-121.95 589.92,-121.95 415.92,-121.95 415.92,-121.95 409.92,-121.95 403.92,-115.95 403.92,-109.95 403.92,-109.95 403.92,-20 403.92,-20 403.92,-14 409.92,-8 415.92,-8\"/>\n",
       "<text text-anchor=\"middle\" x=\"588.92\" y=\"-15.8\" font-family=\"Times,serif\" font-size=\"14.00\">4</text>\n",
       "</g>\n",
       "<!-- scalar (support) -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>scalar (support)</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"91.92\" cy=\"-76.48\" rx=\"91.85\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"91.92\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">scalar (support)</text>\n",
       "<text text-anchor=\"middle\" x=\"91.92\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"91.92\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
       "</g>\n",
       "<!-- vector (implicit) -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>vector (implicit)</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"294.92\" cy=\"-76.48\" rx=\"92.76\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"294.92\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">vector (implicit)</text>\n",
       "<text text-anchor=\"middle\" x=\"294.92\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"294.92\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
       "</g>\n",
       "<!-- vector (explicit) -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>vector (explicit)</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"502.92\" cy=\"-76.48\" rx=\"91.43\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"502.92\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">vector (explicit)</text>\n",
       "<text text-anchor=\"middle\" x=\"502.92\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"502.92\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f642591f250>"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with pm.Model() as pmodel:\n",
    "    pm.Normal(\"scalar (support)\")\n",
    "    pm.Normal(\"vector (implicit)\", mu=[1, 2, 3])\n",
    "    pm.Normal(\"vector (explicit)\", shape=(4,))\n",
    "\n",
    "pm.model_to_graphviz(pmodel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Dims"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyMC supports the concept of `dims`. With many random variables it can become confusing which dimensionality corresponds to which \"real world\" idea, e.g. number of observations, number of treated units etc. The `dims` argument is an additional human-readable label that can convey this meaning. When used alone, `dims` must be combined with explicit `shape` information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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 7.1.0 (20230122.1345)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"138pt\" height=\"138pt\"\n",
       " viewBox=\"0.00 0.00 138.00 137.95\" 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 133.95)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-133.95 134,-133.95 134,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>clusteryear (3)</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M20,-8C20,-8 110,-8 110,-8 116,-8 122,-14 122,-20 122,-20 122,-109.95 122,-109.95 122,-115.95 116,-121.95 110,-121.95 110,-121.95 20,-121.95 20,-121.95 14,-121.95 8,-115.95 8,-109.95 8,-109.95 8,-20 8,-20 8,-14 14,-8 20,-8\"/>\n",
       "<text text-anchor=\"middle\" x=\"86\" y=\"-15.8\" font-family=\"Times,serif\" font-size=\"14.00\">year (3)</text>\n",
       "</g>\n",
       "<!-- profit -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>profit</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"65\" cy=\"-76.48\" rx=\"49.49\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">profit</text>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f6425ac2b60>"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with pm.Model() as model:\n",
    "    pm.Normal(\"profit\", shape=(3,), dims=\"year\")\n",
    "\n",
    "pm.model_to_graphviz(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Where `dims` can become increasingly powerful is with the use of `coords` specified in the model itself. This gives a unique label to each `dim` entry, rendering it much more meaningful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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 7.1.0 (20230122.1345)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"138pt\" height=\"138pt\"\n",
       " viewBox=\"0.00 0.00 138.00 137.95\" 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 133.95)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-133.95 134,-133.95 134,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>clusteryear (3)</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M20,-8C20,-8 110,-8 110,-8 116,-8 122,-14 122,-20 122,-20 122,-109.95 122,-109.95 122,-115.95 116,-121.95 110,-121.95 110,-121.95 20,-121.95 20,-121.95 14,-121.95 8,-115.95 8,-109.95 8,-109.95 8,-20 8,-20 8,-14 14,-8 20,-8\"/>\n",
       "<text text-anchor=\"middle\" x=\"86\" y=\"-15.8\" font-family=\"Times,serif\" font-size=\"14.00\">year (3)</text>\n",
       "</g>\n",
       "<!-- profit -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>profit</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"65\" cy=\"-76.48\" rx=\"49.49\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">profit</text>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"65\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f6425ac3c40>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coords = {\"year\": [2020, 2021, 2022]}\n",
    "with pm.Model(coords=coords) as model:\n",
    "    pm.Normal(\"profit\", dims=\"year\")\n",
    "\n",
    "pm.model_to_graphviz(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In this case the dimensionality of the distribution can actually defined by the `dims` used. We did not have to pass shape or define implicit batched dimensions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let us to review the different dimensionality flavours with a Multivariate Normal example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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 7.1.0 (20230122.1345)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"586pt\" height=\"138pt\"\n",
       " viewBox=\"0.00 0.00 586.00 137.95\" 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 133.95)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-133.95 582,-133.95 582,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>clustersupport (3)</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M20,-8C20,-8 142,-8 142,-8 148,-8 154,-14 154,-20 154,-20 154,-109.95 154,-109.95 154,-115.95 148,-121.95 142,-121.95 142,-121.95 20,-121.95 20,-121.95 14,-121.95 8,-115.95 8,-109.95 8,-109.95 8,-20 8,-20 8,-14 14,-8 20,-8\"/>\n",
       "<text text-anchor=\"middle\" x=\"106\" y=\"-15.8\" font-family=\"Times,serif\" font-size=\"14.00\">support (3)</text>\n",
       "</g>\n",
       "<g id=\"clust2\" class=\"cluster\">\n",
       "<title>clusterbatch (4) x support (3)</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M174,-8C174,-8 558,-8 558,-8 564,-8 570,-14 570,-20 570,-20 570,-109.95 570,-109.95 570,-115.95 564,-121.95 558,-121.95 558,-121.95 174,-121.95 174,-121.95 168,-121.95 162,-115.95 162,-109.95 162,-109.95 162,-20 162,-20 162,-14 168,-8 174,-8\"/>\n",
       "<text text-anchor=\"middle\" x=\"481.5\" y=\"-15.8\" font-family=\"Times,serif\" font-size=\"14.00\">batch (4) x support (3)</text>\n",
       "</g>\n",
       "<!-- vector -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>vector</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"81\" cy=\"-76.48\" rx=\"65.11\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"81\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">vector</text>\n",
       "<text text-anchor=\"middle\" x=\"81\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"81\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">MvNormal</text>\n",
       "</g>\n",
       "<!-- matrix (explicit) -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>matrix (explicit)</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"469\" cy=\"-76.48\" rx=\"92.76\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"469\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">matrix (explicit)</text>\n",
       "<text text-anchor=\"middle\" x=\"469\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"469\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">MvNormal</text>\n",
       "</g>\n",
       "<!-- matrix (implicit) -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>matrix (implicit)</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"264\" cy=\"-76.48\" rx=\"94.09\" ry=\"37.45\"/>\n",
       "<text text-anchor=\"middle\" x=\"264\" y=\"-87.78\" font-family=\"Times,serif\" font-size=\"14.00\">matrix (implicit)</text>\n",
       "<text text-anchor=\"middle\" x=\"264\" y=\"-72.78\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
       "<text text-anchor=\"middle\" x=\"264\" y=\"-57.78\" font-family=\"Times,serif\" font-size=\"14.00\">MvNormal</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7f6425711510>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coords = {\n",
    "    \"batch\": [0, 1, 2, 3],\n",
    "    \"support\": [\"A\", \"B\", \"C\"],\n",
    "}\n",
    "with pm.Model(coords=coords) as model:\n",
    "    pm.MvNormal(\"vector\", mu=[0, 0, 0], cov=np.eye(3), dims=(\"support\",))\n",
    "    pm.MvNormal(\"matrix (implicit)\", mu=np.zeros((4, 3)), cov=np.eye(3), dims=(\"batch\", \"support\"))\n",
    "    pm.MvNormal(\n",
    "        \"matrix (explicit)\", mu=[0, 0, 0], cov=np.eye(3), shape=(4, 3), dims=(\"batch\", \"support\")\n",
    "    )\n",
    "\n",
    "pm.model_to_graphviz(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip} For final model publication we suggest dims and coords as the labels will be passed to {class}`arviz.InferenceData`. This is both best practice transparency and readability for others. It also is useful in single developer workflows, for example, in cases where there is a 3 dimensional or higher distribution it'll help indicate which dimension corresponds to which model concept.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Tips for debugging shape issues"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "While we provide all these tools for convenience, and while PyMC does it best to understand user intent, the result of mixed dimensionality tools may not always result in the final dimensionality intended. Sometimes the model may not indicate an error until sampling, or not indicate an issue at all. When working with dimensionality, particular more complex ones we suggest:\n",
    "\n",
    "* Using `model_to_graphviz` to visualize your model before sampling\n",
    "* Using `draw` or `sample_prior predictive` to catch errors early\n",
    "* Inspecting the returned `az.InferenceData` object to ensure all array sizes are as intended\n",
    "* Defining shapes with prime numbers when tracking down errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last updated: Mon Jul 17 2023\n",
      "\n",
      "Python implementation: CPython\n",
      "Python version       : 3.10.9\n",
      "IPython version      : 8.11.0\n",
      "\n",
      "pytensor: 2.12.3\n",
      "\n",
      "pymc    : 5.2.0\n",
      "numpy   : None\n",
      "pytensor: 2.12.3\n",
      "\n",
      "Watermark: 2.3.1\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "\n",
    "%watermark -n -u -v -iv -w -p pytensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "interpreter": {
   "hash": "f574fac5b7e4a41f7640949d1e1759089329dd116ff7b389caa9cf21f93d872d"
  },
  "kernelspec": {
   "display_name": "pymc",
   "language": "python",
   "name": "pymc"
  },
  "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.10.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}