Update to PLDI artifact.

Author: femtomc

experiments.ipynb

@@ -54,7 +54,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "0a82fac8-f6c1-454b-bbf5-2a9274612c1b",
    "metadata": {},
    "outputs": [
@@ -65,6 +65,22 @@
       "/home/femtomc/.cache/pypoetry/virtualenvs/programmable-vi-pldi-2024-wktT2A4B-py3.10/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
       "  from .autonotebook import tqdm as notebook_tqdm\n"
      ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "GenJAX VI timings:\n",
+      "Batch sizes: [64, 128, 256, 512, 1024]\n",
+      "(array([ 3.2510662,  5.572334 , 10.231137 , 17.31219  , 31.963547 ],\n",
+      "      dtype=float32), array([0.21272574, 2.0414262 , 2.1288395 , 1.691629  , 5.045877  ],\n",
+      "      dtype=float32))\n",
+      "Handcoded timings:\n",
+      "Batch sizes: [64, 128, 256, 512, 1024]\n",
+      "(array([ 3.345306 ,  5.3966055,  9.804233 , 15.3170595, 31.775038 ],\n",
+      "      dtype=float32), array([0.06193171, 0.04953469, 2.064177  , 4.5305977 , 6.054197  ],\n",
+      "      dtype=float32))\n"
+     ]
     }
    ],
    "source": [
@@ -100,7 +116,23 @@
    "execution_count": null,
    "id": "27525362-6d68-4d8a-841a-ec0c3011e756",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[1mpoetry run pytest experiments/table_2_benchmark_timings --benchmark-disable-gc\u001b[0m\n",
+      "\u001b[1m============================= test session starts ==============================\u001b[0m\n",
+      "platform linux -- Python 3.10.13, pytest-8.0.2, pluggy-1.4.0\n",
+      "benchmark: 4.0.0 (defaults: timer=time.perf_counter disable_gc=True min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)\n",
+      "rootdir: /home/femtomc/programmable-vi-pldi-2024\n",
+      "plugins: jaxtyping-0.2.28, benchmark-4.0.0, anyio-4.3.0, typeguard-2.13.3\n",
+      "collected 6 items                                                              \u001b[0m\u001b[1m\u001b[1m\u001b[1m\n",
+      "\n",
+      "experiments/table_2_benchmark_timings/test_genjax_enum_air_benchmark.py "
+     ]
+    }
+   ],
    "source": [
     "! just table_2"
    ]

experiments/fig_7_air_estimator_evaluation/genjax_rws_air.py

@@ -670,7 +670,7 @@ def draw_many(imgs, zs, title):
 evaluate_accuracy = count_accuracy(mnist, true_counts, guide, batch_size=1000)
 
 
-def train(key, n=1, num_epochs=40, batch_size=64, learning_rate=1.0e-3):
+def train(key, n=2, num_epochs=40, batch_size=64, learning_rate=1.0e-3):
     def svi_update(model, guide, optimiser):
         def batch_updater(key, params, opt_state, data_batch):
             def p_grads(key, params, data):
@@ -788,7 +788,7 @@ def body_fn(carry, xs):
 
 key, sub_key = jax.random.split(key)
 (p_losses, q_losses), accuracy, wall_clock_times, params = train(
-    sub_key, learning_rate=1.0e-3, n=10, batch_size=64, num_epochs=40
+    sub_key, learning_rate=3.0e-3, n=10, batch_size=64, num_epochs=40
 )
 
 arr = np.array([p_losses, q_losses, accuracy, wall_clock_times])