Add figure to readme

Author: smsharma

README.md

@@ -5,6 +5,8 @@ Carolina Cuesta-Lazaro and Siddharth Mishra-Sharma
 [![License: MIT](https://img.shields.io/badge/License-MIT-red.svg)](https://opensource.org/licenses/MIT)
 [![arXiv](https://img.shields.io/badge/arXiv-2311.xxxxx%20-green.svg)](https://arxiv.org/abs/2311.xxxxx)
 
+![Figure.](notebooks/plots/diffusion_fig.png)
+
 
 ## Abstract
 

notebooks/04_viz_pos.ipynb

@@ -269,26 +269,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 35,
+   "execution_count": 36,
    "metadata": {},
-   "outputs": [
-    {
-     "ename": "TypeError",
-     "evalue": "mul got incompatible shapes for broadcasting: (10,), (3,).",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[35], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnorm_dict\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mstd\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m \u001b[38;5;241m+\u001b[39m norm_dict[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmean\u001b[39m\u001b[38;5;124m'\u001b[39m]\n",
-      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:256\u001b[0m, in \u001b[0;36m_defer_to_unrecognized_arg.<locals>.deferring_binary_op\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m    254\u001b[0m args \u001b[38;5;241m=\u001b[39m (other, \u001b[38;5;28mself\u001b[39m) \u001b[38;5;28;01mif\u001b[39;00m swap \u001b[38;5;28;01melse\u001b[39;00m (\u001b[38;5;28mself\u001b[39m, other)\n\u001b[1;32m    255\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(other, _accepted_binop_types):\n\u001b[0;32m--> 256\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mbinary_op\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    257\u001b[0m \u001b[38;5;66;03m# Note: don't use isinstance here, because we don't want to raise for\u001b[39;00m\n\u001b[1;32m    258\u001b[0m \u001b[38;5;66;03m# subclasses, e.g. NamedTuple objects that may override operators.\u001b[39;00m\n\u001b[1;32m    259\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(other) \u001b[38;5;129;01min\u001b[39;00m _rejected_binop_types:\n",
-      "    \u001b[0;31m[... skipping hidden 12 frame]\u001b[0m\n",
-      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/jax/_src/numpy/ufuncs.py:97\u001b[0m, in \u001b[0;36m_maybe_bool_binop.<locals>.fn\u001b[0;34m(x1, x2)\u001b[0m\n\u001b[1;32m     95\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfn\u001b[39m(x1, x2, \u001b[38;5;241m/\u001b[39m):\n\u001b[1;32m     96\u001b[0m   x1, x2 \u001b[38;5;241m=\u001b[39m promote_args(numpy_fn\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, x1, x2)\n\u001b[0;32m---> 97\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mlax_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx1\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx2\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mif\u001b[39;00m x1\u001b[38;5;241m.\u001b[39mdtype \u001b[38;5;241m!=\u001b[39m np\u001b[38;5;241m.\u001b[39mbool_ \u001b[38;5;28;01melse\u001b[39;00m bool_lax_fn(x1, x2)\n",
-      "    \u001b[0;31m[... skipping hidden 7 frame]\u001b[0m\n",
-      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/jax/_src/lax/lax.py:1591\u001b[0m, in \u001b[0;36mbroadcasting_shape_rule\u001b[0;34m(name, *avals)\u001b[0m\n\u001b[1;32m   1589\u001b[0m       result_shape\u001b[38;5;241m.\u001b[39mappend(non_1s[\u001b[38;5;241m0\u001b[39m])\n\u001b[1;32m   1590\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1591\u001b[0m       \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m got incompatible shapes for broadcasting: \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m   1592\u001b[0m                       \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;28mmap\u001b[39m(\u001b[38;5;28mstr\u001b[39m,\u001b[38;5;250m \u001b[39m\u001b[38;5;28mmap\u001b[39m(\u001b[38;5;28mtuple\u001b[39m,\u001b[38;5;250m \u001b[39mshapes)))\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m   1594\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtuple\u001b[39m(result_shape)\n",
-      "\u001b[0;31mTypeError\u001b[0m: mul got incompatible shapes for broadcasting: (10,), (3,)."
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "x = x * norm_dict['std'] + norm_dict['mean']"
    ]