diff --git a/notebooks/inference-testing.ipynb b/notebooks/inference-testing.ipynb
index efb7a263..3784d6e4 100644
--- a/notebooks/inference-testing.ipynb
+++ b/notebooks/inference-testing.ipynb
@@ -492,132 +492,132 @@
        "  stroke-width: 0.8px;\n",
        "}\n",
        "</style><pre class='xr-text-repr-fallback'>&lt;xarray.DataArray (x: 100, y: 100, t: 10)&gt; Size: 800kB\n",
-       "array([[[0.44893786, 0.42089057, 0.62811935, ..., 0.77825309,\n",
-       "         0.71424518, 0.27544911],\n",
-       "        [0.18512313, 0.22459034, 0.32941286, ..., 0.72642438,\n",
-       "         0.453367  , 0.79092317],\n",
-       "        [0.75001122, 0.75201415, 0.49362845, ..., 0.12351849,\n",
-       "         0.07800031, 0.5129466 ],\n",
+       "array([[[0.80102402, 0.31982866, 0.25465188, ..., 0.31021283,\n",
+       "         0.03268302, 0.7941252 ],\n",
+       "        [0.4382327 , 0.60519493, 0.86055512, ..., 0.91560426,\n",
+       "         0.75151732, 0.045112  ],\n",
+       "        [0.56324046, 0.16672896, 0.95892792, ..., 0.75294859,\n",
+       "         0.62661621, 0.9507573 ],\n",
        "        ...,\n",
-       "        [0.79767373, 0.04314154, 0.42138962, ..., 0.64572546,\n",
-       "         0.68486366, 0.26731422],\n",
-       "        [0.55990821, 0.66252347, 0.35829701, ..., 0.2245369 ,\n",
-       "         0.210521  , 0.3942034 ],\n",
-       "        [0.26561369, 0.32802417, 0.33247419, ..., 0.89834246,\n",
-       "         0.85079973, 0.30808448]],\n",
-       "\n",
-       "       [[0.49745507, 0.47563665, 0.25939342, ..., 0.95437792,\n",
-       "         0.3529035 , 0.22432583],\n",
-       "        [0.93810888, 0.36236317, 0.90214027, ..., 0.52733391,\n",
-       "         0.08590167, 0.39300478],\n",
-       "        [0.39701346, 0.87476725, 0.07576221, ..., 0.6206168 ,\n",
-       "         0.590455  , 0.18998422],\n",
+       "        [0.8942051 , 0.67152092, 0.5936959 , ..., 0.88955481,\n",
+       "         0.00104593, 0.25334781],\n",
+       "        [0.20207265, 0.89166716, 0.11930909, ..., 0.44454162,\n",
+       "         0.16769459, 0.13463857],\n",
+       "        [0.52366905, 0.77434724, 0.18922243, ..., 0.60982095,\n",
+       "         0.81261007, 0.35591865]],\n",
+       "\n",
+       "       [[0.36699177, 0.29720945, 0.47093257, ..., 0.08379803,\n",
+       "         0.52233317, 0.62779805],\n",
+       "        [0.78408776, 0.29027363, 0.06707094, ..., 0.83199475,\n",
+       "         0.26807477, 0.93323152],\n",
+       "        [0.70058356, 0.83639276, 0.82536327, ..., 0.94036962,\n",
+       "         0.79252184, 0.83713594],\n",
        "...\n",
-       "        [0.59384608, 0.60548133, 0.86872946, ..., 0.18636072,\n",
-       "         0.23673955, 0.47465824],\n",
-       "        [0.53055158, 0.63400677, 0.74856164, ..., 0.57724106,\n",
-       "         0.12961194, 0.23441086],\n",
-       "        [0.04898592, 0.96045848, 0.96103484, ..., 0.11542319,\n",
-       "         0.03916002, 0.73413531]],\n",
-       "\n",
-       "       [[0.71738595, 0.97645637, 0.98795853, ..., 0.27255291,\n",
-       "         0.01424937, 0.08837012],\n",
-       "        [0.64385468, 0.25559843, 0.55311384, ..., 0.45737679,\n",
-       "         0.6211287 , 0.40359951],\n",
-       "        [0.42621607, 0.309378  , 0.1063637 , ..., 0.23134538,\n",
-       "         0.04260683, 0.64022588],\n",
+       "        [0.39766555, 0.35994714, 0.50271116, ..., 0.46796155,\n",
+       "         0.97140383, 0.51062791],\n",
+       "        [0.70055487, 0.81950773, 0.76781762, ..., 0.05735074,\n",
+       "         0.73039909, 0.17101624],\n",
+       "        [0.8547392 , 0.21515083, 0.40187337, ..., 0.4583594 ,\n",
+       "         0.53213994, 0.15382577]],\n",
+       "\n",
+       "       [[0.6465284 , 0.88463142, 0.33543765, ..., 0.49702278,\n",
+       "         0.40990036, 0.42300881],\n",
+       "        [0.20363819, 0.02142565, 0.66128113, ..., 0.56253861,\n",
+       "         0.33799959, 0.16641441],\n",
+       "        [0.67934574, 0.38965227, 0.68431764, ..., 0.56820253,\n",
+       "         0.18640022, 0.111049  ],\n",
        "        ...,\n",
-       "        [0.8286403 , 0.18384525, 0.80661243, ..., 0.78306387,\n",
-       "         0.29219244, 0.22387774],\n",
-       "        [0.65402132, 0.56136263, 0.86077124, ..., 0.3128178 ,\n",
-       "         0.64353248, 0.52959765],\n",
-       "        [0.76868123, 0.67687677, 0.42182969, ..., 0.86638394,\n",
-       "         0.68939672, 0.14375773]]], shape=(100, 100, 10))\n",
-       "Dimensions without coordinates: x, y, t</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.DataArray</div><div class='xr-array-name'></div><ul class='xr-dim-list'><li><span>x</span>: 100</li><li><span>y</span>: 100</li><li><span>t</span>: 10</li></ul></div><ul class='xr-sections'><li class='xr-section-item'><div class='xr-array-wrap'><input id='section-3955e21b-9898-47d3-90c5-6807f8cbc064' class='xr-array-in' type='checkbox' checked><label for='section-3955e21b-9898-47d3-90c5-6807f8cbc064' title='Show/hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-array-preview xr-preview'><span>0.4489 0.4209 0.6281 0.994 0.4309 ... 0.9977 0.8664 0.6894 0.1438</span></div><div class='xr-array-data'><pre>array([[[0.44893786, 0.42089057, 0.62811935, ..., 0.77825309,\n",
-       "         0.71424518, 0.27544911],\n",
-       "        [0.18512313, 0.22459034, 0.32941286, ..., 0.72642438,\n",
-       "         0.453367  , 0.79092317],\n",
-       "        [0.75001122, 0.75201415, 0.49362845, ..., 0.12351849,\n",
-       "         0.07800031, 0.5129466 ],\n",
+       "        [0.91296116, 0.6995538 , 0.50325961, ..., 0.40427649,\n",
+       "         0.14176453, 0.26158606],\n",
+       "        [0.5348746 , 0.67937558, 0.22366613, ..., 0.48889003,\n",
+       "         0.50030053, 0.1855764 ],\n",
+       "        [0.8714776 , 0.9552773 , 0.06885801, ..., 0.95366137,\n",
+       "         0.73931089, 0.90909682]]], shape=(100, 100, 10))\n",
+       "Dimensions without coordinates: x, y, t</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.DataArray</div><div class='xr-array-name'></div><ul class='xr-dim-list'><li><span>x</span>: 100</li><li><span>y</span>: 100</li><li><span>t</span>: 10</li></ul></div><ul class='xr-sections'><li class='xr-section-item'><div class='xr-array-wrap'><input id='section-a3760498-07be-4ecb-a657-48ddbc638af0' class='xr-array-in' type='checkbox' checked><label for='section-a3760498-07be-4ecb-a657-48ddbc638af0' title='Show/hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-array-preview xr-preview'><span>0.801 0.3198 0.2547 0.8238 0.4707 ... 0.001633 0.9537 0.7393 0.9091</span></div><div class='xr-array-data'><pre>array([[[0.80102402, 0.31982866, 0.25465188, ..., 0.31021283,\n",
+       "         0.03268302, 0.7941252 ],\n",
+       "        [0.4382327 , 0.60519493, 0.86055512, ..., 0.91560426,\n",
+       "         0.75151732, 0.045112  ],\n",
+       "        [0.56324046, 0.16672896, 0.95892792, ..., 0.75294859,\n",
+       "         0.62661621, 0.9507573 ],\n",
        "        ...,\n",
-       "        [0.79767373, 0.04314154, 0.42138962, ..., 0.64572546,\n",
-       "         0.68486366, 0.26731422],\n",
-       "        [0.55990821, 0.66252347, 0.35829701, ..., 0.2245369 ,\n",
-       "         0.210521  , 0.3942034 ],\n",
-       "        [0.26561369, 0.32802417, 0.33247419, ..., 0.89834246,\n",
-       "         0.85079973, 0.30808448]],\n",
-       "\n",
-       "       [[0.49745507, 0.47563665, 0.25939342, ..., 0.95437792,\n",
-       "         0.3529035 , 0.22432583],\n",
-       "        [0.93810888, 0.36236317, 0.90214027, ..., 0.52733391,\n",
-       "         0.08590167, 0.39300478],\n",
-       "        [0.39701346, 0.87476725, 0.07576221, ..., 0.6206168 ,\n",
-       "         0.590455  , 0.18998422],\n",
+       "        [0.8942051 , 0.67152092, 0.5936959 , ..., 0.88955481,\n",
+       "         0.00104593, 0.25334781],\n",
+       "        [0.20207265, 0.89166716, 0.11930909, ..., 0.44454162,\n",
+       "         0.16769459, 0.13463857],\n",
+       "        [0.52366905, 0.77434724, 0.18922243, ..., 0.60982095,\n",
+       "         0.81261007, 0.35591865]],\n",
+       "\n",
+       "       [[0.36699177, 0.29720945, 0.47093257, ..., 0.08379803,\n",
+       "         0.52233317, 0.62779805],\n",
+       "        [0.78408776, 0.29027363, 0.06707094, ..., 0.83199475,\n",
+       "         0.26807477, 0.93323152],\n",
+       "        [0.70058356, 0.83639276, 0.82536327, ..., 0.94036962,\n",
+       "         0.79252184, 0.83713594],\n",
        "...\n",
-       "        [0.59384608, 0.60548133, 0.86872946, ..., 0.18636072,\n",
-       "         0.23673955, 0.47465824],\n",
-       "        [0.53055158, 0.63400677, 0.74856164, ..., 0.57724106,\n",
-       "         0.12961194, 0.23441086],\n",
-       "        [0.04898592, 0.96045848, 0.96103484, ..., 0.11542319,\n",
-       "         0.03916002, 0.73413531]],\n",
-       "\n",
-       "       [[0.71738595, 0.97645637, 0.98795853, ..., 0.27255291,\n",
-       "         0.01424937, 0.08837012],\n",
-       "        [0.64385468, 0.25559843, 0.55311384, ..., 0.45737679,\n",
-       "         0.6211287 , 0.40359951],\n",
-       "        [0.42621607, 0.309378  , 0.1063637 , ..., 0.23134538,\n",
-       "         0.04260683, 0.64022588],\n",
+       "        [0.39766555, 0.35994714, 0.50271116, ..., 0.46796155,\n",
+       "         0.97140383, 0.51062791],\n",
+       "        [0.70055487, 0.81950773, 0.76781762, ..., 0.05735074,\n",
+       "         0.73039909, 0.17101624],\n",
+       "        [0.8547392 , 0.21515083, 0.40187337, ..., 0.4583594 ,\n",
+       "         0.53213994, 0.15382577]],\n",
+       "\n",
+       "       [[0.6465284 , 0.88463142, 0.33543765, ..., 0.49702278,\n",
+       "         0.40990036, 0.42300881],\n",
+       "        [0.20363819, 0.02142565, 0.66128113, ..., 0.56253861,\n",
+       "         0.33799959, 0.16641441],\n",
+       "        [0.67934574, 0.38965227, 0.68431764, ..., 0.56820253,\n",
+       "         0.18640022, 0.111049  ],\n",
        "        ...,\n",
-       "        [0.8286403 , 0.18384525, 0.80661243, ..., 0.78306387,\n",
-       "         0.29219244, 0.22387774],\n",
-       "        [0.65402132, 0.56136263, 0.86077124, ..., 0.3128178 ,\n",
-       "         0.64353248, 0.52959765],\n",
-       "        [0.76868123, 0.67687677, 0.42182969, ..., 0.86638394,\n",
-       "         0.68939672, 0.14375773]]], shape=(100, 100, 10))</pre></div></div></li><li class='xr-section-item'><input id='section-2b43c545-42c2-40e8-ba80-8f2dcbfaa8ed' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-2b43c545-42c2-40e8-ba80-8f2dcbfaa8ed' class='xr-section-summary'  title='Expand/collapse section'>Coordinates: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-fd46f49b-7192-44fd-8f61-994f489a3b60' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-fd46f49b-7192-44fd-8f61-994f489a3b60' class='xr-section-summary'  title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-06020ee5-a9dd-4e49-9bd9-f59be3ffa00f' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-06020ee5-a9dd-4e49-9bd9-f59be3ffa00f' class='xr-section-summary'  title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>"
+       "        [0.91296116, 0.6995538 , 0.50325961, ..., 0.40427649,\n",
+       "         0.14176453, 0.26158606],\n",
+       "        [0.5348746 , 0.67937558, 0.22366613, ..., 0.48889003,\n",
+       "         0.50030053, 0.1855764 ],\n",
+       "        [0.8714776 , 0.9552773 , 0.06885801, ..., 0.95366137,\n",
+       "         0.73931089, 0.90909682]]], shape=(100, 100, 10))</pre></div></div></li><li class='xr-section-item'><input id='section-cfb30e7f-3d72-4bff-8128-fd622b7bc3b7' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-cfb30e7f-3d72-4bff-8128-fd622b7bc3b7' class='xr-section-summary'  title='Expand/collapse section'>Coordinates: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-a0b305aa-ceb4-4fab-9357-d518e239019e' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-a0b305aa-ceb4-4fab-9357-d518e239019e' class='xr-section-summary'  title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-a6cc32f8-9160-416f-8332-bf26ed7186e1' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-a6cc32f8-9160-416f-8332-bf26ed7186e1' class='xr-section-summary'  title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>"
       ],
       "text/plain": [
        "<xarray.DataArray (x: 100, y: 100, t: 10)> Size: 800kB\n",
-       "array([[[0.44893786, 0.42089057, 0.62811935, ..., 0.77825309,\n",
-       "         0.71424518, 0.27544911],\n",
-       "        [0.18512313, 0.22459034, 0.32941286, ..., 0.72642438,\n",
-       "         0.453367  , 0.79092317],\n",
-       "        [0.75001122, 0.75201415, 0.49362845, ..., 0.12351849,\n",
-       "         0.07800031, 0.5129466 ],\n",
+       "array([[[0.80102402, 0.31982866, 0.25465188, ..., 0.31021283,\n",
+       "         0.03268302, 0.7941252 ],\n",
+       "        [0.4382327 , 0.60519493, 0.86055512, ..., 0.91560426,\n",
+       "         0.75151732, 0.045112  ],\n",
+       "        [0.56324046, 0.16672896, 0.95892792, ..., 0.75294859,\n",
+       "         0.62661621, 0.9507573 ],\n",
        "        ...,\n",
-       "        [0.79767373, 0.04314154, 0.42138962, ..., 0.64572546,\n",
-       "         0.68486366, 0.26731422],\n",
-       "        [0.55990821, 0.66252347, 0.35829701, ..., 0.2245369 ,\n",
-       "         0.210521  , 0.3942034 ],\n",
-       "        [0.26561369, 0.32802417, 0.33247419, ..., 0.89834246,\n",
-       "         0.85079973, 0.30808448]],\n",
-       "\n",
-       "       [[0.49745507, 0.47563665, 0.25939342, ..., 0.95437792,\n",
-       "         0.3529035 , 0.22432583],\n",
-       "        [0.93810888, 0.36236317, 0.90214027, ..., 0.52733391,\n",
-       "         0.08590167, 0.39300478],\n",
-       "        [0.39701346, 0.87476725, 0.07576221, ..., 0.6206168 ,\n",
-       "         0.590455  , 0.18998422],\n",
+       "        [0.8942051 , 0.67152092, 0.5936959 , ..., 0.88955481,\n",
+       "         0.00104593, 0.25334781],\n",
+       "        [0.20207265, 0.89166716, 0.11930909, ..., 0.44454162,\n",
+       "         0.16769459, 0.13463857],\n",
+       "        [0.52366905, 0.77434724, 0.18922243, ..., 0.60982095,\n",
+       "         0.81261007, 0.35591865]],\n",
+       "\n",
+       "       [[0.36699177, 0.29720945, 0.47093257, ..., 0.08379803,\n",
+       "         0.52233317, 0.62779805],\n",
+       "        [0.78408776, 0.29027363, 0.06707094, ..., 0.83199475,\n",
+       "         0.26807477, 0.93323152],\n",
+       "        [0.70058356, 0.83639276, 0.82536327, ..., 0.94036962,\n",
+       "         0.79252184, 0.83713594],\n",
        "...\n",
-       "        [0.59384608, 0.60548133, 0.86872946, ..., 0.18636072,\n",
-       "         0.23673955, 0.47465824],\n",
-       "        [0.53055158, 0.63400677, 0.74856164, ..., 0.57724106,\n",
-       "         0.12961194, 0.23441086],\n",
-       "        [0.04898592, 0.96045848, 0.96103484, ..., 0.11542319,\n",
-       "         0.03916002, 0.73413531]],\n",
-       "\n",
-       "       [[0.71738595, 0.97645637, 0.98795853, ..., 0.27255291,\n",
-       "         0.01424937, 0.08837012],\n",
-       "        [0.64385468, 0.25559843, 0.55311384, ..., 0.45737679,\n",
-       "         0.6211287 , 0.40359951],\n",
-       "        [0.42621607, 0.309378  , 0.1063637 , ..., 0.23134538,\n",
-       "         0.04260683, 0.64022588],\n",
+       "        [0.39766555, 0.35994714, 0.50271116, ..., 0.46796155,\n",
+       "         0.97140383, 0.51062791],\n",
+       "        [0.70055487, 0.81950773, 0.76781762, ..., 0.05735074,\n",
+       "         0.73039909, 0.17101624],\n",
+       "        [0.8547392 , 0.21515083, 0.40187337, ..., 0.4583594 ,\n",
+       "         0.53213994, 0.15382577]],\n",
+       "\n",
+       "       [[0.6465284 , 0.88463142, 0.33543765, ..., 0.49702278,\n",
+       "         0.40990036, 0.42300881],\n",
+       "        [0.20363819, 0.02142565, 0.66128113, ..., 0.56253861,\n",
+       "         0.33799959, 0.16641441],\n",
+       "        [0.67934574, 0.38965227, 0.68431764, ..., 0.56820253,\n",
+       "         0.18640022, 0.111049  ],\n",
        "        ...,\n",
-       "        [0.8286403 , 0.18384525, 0.80661243, ..., 0.78306387,\n",
-       "         0.29219244, 0.22387774],\n",
-       "        [0.65402132, 0.56136263, 0.86077124, ..., 0.3128178 ,\n",
-       "         0.64353248, 0.52959765],\n",
-       "        [0.76868123, 0.67687677, 0.42182969, ..., 0.86638394,\n",
-       "         0.68939672, 0.14375773]]], shape=(100, 100, 10))\n",
+       "        [0.91296116, 0.6995538 , 0.50325961, ..., 0.40427649,\n",
+       "         0.14176453, 0.26158606],\n",
+       "        [0.5348746 , 0.67937558, 0.22366613, ..., 0.48889003,\n",
+       "         0.50030053, 0.1855764 ],\n",
+       "        [0.8714776 , 0.9552773 , 0.06885801, ..., 0.95366137,\n",
+       "         0.73931089, 0.90909682]]], shape=(100, 100, 10))\n",
        "Dimensions without coordinates: x, y, t"
       ]
      },
@@ -724,27 +724,101 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "(200, 100, 5)\n"
+      "{'y': 100, 'x': 100, 't': 5}\n"
      ]
     }
    ],
    "source": [
-    "out = predict_on_array(\n",
-    "    dataset=ds, \n",
-    "    model=mad, \n",
-    "    output_tensor_dim=dict(y=20, x=10, x_new=5),\n",
-    "    new_dim=[\"x_new\"],\n",
-    "    resample_dim=[\"x\", \"y\"]\n",
+    "out_size_dict = _get_output_array_size(\n",
+    "    bgen = ds.X_generator,\n",
+    "    output_tensor_dim = dict(y=10, x=10, t=5),\n",
+    "    new_dim = [\"t\"],\n",
+    "    resample_dim = [\"y\", \"x\", \"t\"]\n",
     ")\n",
-    "print(out.shape)"
+    "print(out_size_dict)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from functools import partial"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "get_array_size_partial = partial(_get_output_array_size, bgen=ds.X_generator)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_cases = [\n",
+    "    {\n",
+    "        \"name\": \"Same dims and same input sizes 1\",\n",
+    "        \"function_inputs\": {\n",
+    "            \"output_tensor_dim\": dict(y=10, x=10),\n",
+    "            \"new_dim\": [],\n",
+    "            \"resample_dim\": [\"x\", \"y\"]\n",
+    "        },\n",
+    "        \"expected_output\": dict(y=100, x=100)\n",
+    "    },\n",
+    "    {\n",
+    "        \"name\": \"New dim added\",\n",
+    "        \"function_inputs\": {\n",
+    "            \"output_tensor_dim\": dict(y=10, x=10, t=5),\n",
+    "            \"new_dim\": [\"t\"],\n",
+    "            \"resample_dim\": [\"y\", \"x\"]\n",
+    "        },\n",
+    "        \"expected_output\": dict(y=100, x=100, t=5)\n",
+    "    }\n",
+    "]         "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test case 0 passed\n",
+      "Test case 1 passed\n"
+     ]
+    }
+   ],
+   "source": [
+    "for i, test_case in enumerate(test_cases):\n",
+    "    true_output = get_array_size_partial(**test_case[\"function_inputs\"])\n",
+    "    success = true_output == test_case[\"expected_output\"]\n",
+    "    message = \"passed\" if success else \"failed\"\n",
+    "    print(f\"Test case {i} {message}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -763,7 +837,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.16"
+   "version": "3.12.11"
   },
   "nbdime-conflicts": {
    "local_diff": [