Add minibatch proposal

Author: zaxtax

VI_Overview.ipynb

@@ -460,7 +460,7 @@
     "    def __init__(self, model=None, optimizers=None):\n",
     "        ...\n",
     "\n",
-    "    def step(self, batched_data):\n",
+    "    def step(self, batch):\n",
     "        ...\n",
     "        return loss\n",
     "````\n",
@@ -506,6 +506,7 @@
     "\n",
     "````python\n",
     "with pm.Model() as model:\n",
+    "    data = pm.Data(\"data\", ...)\n",
     "    x = pm.Normal(\"x\", 0, 1)\n",
     "    y = pm.Normal(\"y\", x, 1, observed=data)\n",
     "\n",
@@ -530,10 +531,44 @@
     "````"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "7f97c341-e9bb-4301-b452-d006d6408cec",
+   "metadata": {},
+   "source": [
+    "### Reworking Minibatch\n",
+    "\n",
+    "Another small change we should consider is moving `pm.Minibatch` out of the model. Max already has a [proposal](https://github.com/pymc-devs/pymc/issues/7496) that I think can be adopted with only a few changes.\n",
+    "\n",
+    "I think where before we explicitly minibatch the data, instead we have dataloaders that stream in updates to the model.\n",
+    "\n",
+    "````python\n",
+    "with pm.Model() as model:\n",
+    "    data = pm.Data(\"data\", None)\n",
+    "    x = pm.Normal(\"x\", 0, 1)\n",
+    "    y = pm.Normal(\"y\", x, 1, observed=data)\n",
+    "\n",
+    "dataloader = pm.Dataloader(np.random.normal(10_000, 2), batch_size=64)\n",
+    "\n",
+    "with model:\n",
+    "    trainer = Trainer(method=ADVI(), dataloader=dataloader)\n",
+    "    trainer.fit(n=10_000)\n",
+    "````\n",
+    "\n",
+    "Importantly, the model doesn't need to know about the dataloader. We will need to tweak the inference object, but it's not so bad.\n",
+    "\n",
+    "````python\n",
+    "class ADVI(Inference):\n",
+    "    def step(self, batch):\n",
+    "        self.model.set_data(\"data\", batch)\n",
+    "        ...\n",
+    "````"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "50fcb3a1-4467-4ace-acdd-666e4f342984",
+   "id": "220ba769-fb8f-47a7-82b6-ab6ca13ad61e",
    "metadata": {},
    "outputs": [],
    "source": []