Initial commit

Co-authored-by: Aradhana Dube <a.dube@rug.nl> Co-authored-by: Renzo I. Barraza Altamirano <r.i.barraza.altamirano@rug.nl> Co-authored-by: Paolo Gibertini <p.gibertini@rug.nl> Co-authored-by: Luca D. Fehlings <l.d.fehlings@rug.nl>
2026-09-30 21:03:06 +02:00 · 2026-02-27 17:43:31 +01:00
commit 9fabbdefc0
@@ -0,0 +1,158 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "437b0062",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import pandas as pd\n",
+    "import seaborn as sns"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "20c6c596",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA1AAAAEmCAYAAABs/3HEAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAPYQAAD2EBqD+naQAA8y5JREFUeJzsnXd8W9X5/99X23vPxFlk7xAghJmywixhl5bShNGWJt/CL9BBBxRKCx2MQilQSggdFAqFQMsmECgESCBk7+kkXvG2ZVvz/v44d0mWbdmJLTs+79fLL8lXV7pH0pV0Pud5ns+jqKqqIpFIJBKJRCKRSCSSLrElegASiUQikUgkEolEMlCQAkoikUgkEolEIpFI4kQKKIlEIpFIJBKJRCKJEymgJBKJRCKRSCQSiSROpICSSCQSiUQikUgkkjiRAkoikUgkEolEIpFI4kQKKIlEIpFIJBKJRCKJEymgJBKJRCKRSCQSiSROHIkeQF8TDocpKysjLS0NRVESPRyJRCIZVKiqSlNTE8XFxdhscg1PR/42SSQSSWLoye/SoBNQZWVllJSUJHoYEolEMqjZv38/Q4cOTfQw+g3yt0kikUgSS3d+lwadgEpLSwPEi5Senp7g0UgkEsngorGxkZKSEuO7WCKQv00SiUSSGHryuzToBJSeGpGeni5/pCQSiSRByDS1SORvk0QikSSW7vwuyQR0iUQikUgkEolEIokTKaAkEolEIpFIJBKJJE6kgJJIJBKJRCKRSCSSOBl0NVASiUQikUgk/QVVVQkGg4RCoUQPZdBjt9txOByyRlPSJVJASSQSiUQikSQAv99PeXk5LS0tiR6KRCM5OZmioiJcLleihyLpx0gBJZFIJBLJAEdVVZp8QdI9zkQPRRIn4XCYPXv2YLfbKS4uxuVyychHAlFVFb/fz6FDh9izZw9jxoyRzb4lHSIFlEQiMWmqhFcWwsz5MOHCRI9GIpHEQX2Ln9teWE99i5/nvn0iDruc9A0E/H4/4XCYkpISkpOTEz0cCZCUlITT6WTfvn34/X48Hk+ihyTpp8hvWYlEYrLzHfH32eOJHolEIomThtYAu3fv4Mt91dz/zvZED0fSTWSUo38h3w9JPMizRCKRmPiaxWVzVWLHIZFI4mZ4pouXMx/iZdcdfPDBct7fJj+/EolE0ptIASWRSEwCXnHplRMwiWTAcGgrGb5yptr28KrrZ+z9522UV9ckelQSiURy1CIFlEQiMfFrTlCtdRAKJHYsEokkPgqnwMLVhCbOw6GEWcArKI+dRPDgukSPTDLIWbFiBYqiUF9fH/d9RowYwUMPPdRrY5JIjgRSQEkkEhO/17zurU7cOCQSSfdIK8B+5TNUXbiUCjWHwlAFypK5sPnVRI9M0o+ZP38+iqLw3e9+t91tCxcuRFEU5s+f3/cDk0j6OVJASSQSk4BVQB1K3DgkEkmPyD/uEv4wdikfhqZgD7XCv74JH/wOVDXRQ5P0U0pKSnjuuedobW01trW1tfHss88ybNiwBI5MIum/9AsB9eijjzJixAg8Hg+zZs1i1apVHe47Z84cFEVp93fBBRf04YglkqMUv6WZo6yDkkgGJCVDilkQ+CEfZF0mNrx/D/zn+xAOJXZgkn7JscceS0lJCS+99JKx7aWXXmLYsGHMmDHD2Obz+fj+979Pfn4+Ho+HU045hdWrV0c81uuvv87YsWNJSkriK1/5Cnv37m13vI8++ohTTz2VpKQkSkpK+P73v4/X6223n0TSn0m4gHr++edZvHgxd955J2vWrGHatGnMnTuXqqrYk7eXXnqJ8vJy42/jxo3Y7XauuOKKPh65RHIUYk3ha5YRKIlkIDKuII0Qdu5V58OFD4FigzV/hZe/A6Fgoocn6Ydcd911PP3008b/S5YsYcGCBRH7/PCHP+Tf//43zzzzDGvWrGH06NHMnTuX2tpaAPbv38+ll17KRRddxNq1a7nhhhv48Y9/HPEYu3bt4txzz+Wyyy5j/fr1PP/883z00UcsWrSo95+kRHIESbiAeuCBB7jxxhtZsGABEydO5PHHHyc5OZklS5bE3D87O5vCwkLj75133iE5OVkKKInkSCBT+CSSAc/YgjQAdh/yEpjxLbh8CdgcsOEFeHE+BP2JHaCk33HNNdfw0UcfsW/fPvbt28fHH3/MNddcY9zu9Xp57LHH+N3vfsd5553HxIkTefLJJ0lKSuKpp54C4LHHHuOYY47h/vvvZ9y4cXzjG99oVz9177338o1vfINbbrmFMWPGcNJJJ/Hwww/z17/+lba2tr58yhLJYeFI5MH9fj9ffPEFt99+u7HNZrNx1lln8cknn8T1GE899RRf+9rXSElJiXm7z+fD5/MZ/zc2Nh7eoCWSoxmZwieRDHiGZCaR4rLj9YfYV+Nl9KRLwOGBf10LW/4DL38bLn8aFCXRQ5X0E/Ly8rjgggtYunQpqqpywQUXkJuba9y+a9cuAoEAJ598srHN6XRywgknsGXLFgC2bNnCrFmzIh539uzZEf+vW7eO9evX849//MPYpqoq4XCYPXv2MGHChN54ehLJESehEajq6mpCoRAFBQUR2wsKCqioqOjy/qtWrWLjxo3ccMMNHe5z7733kpGRYfyVlJQc9rglkqMWmcInkQx4bDaFMVoUaluF1hx73Hlw9XNgc8Kml+HzpxI4Qkl/5LrrrmPp0qU888wzXHfddb1yjObmZr7zne+wdu1a42/dunXs2LGDY445pleOKZH0BglP4TscnnrqKaZMmcIJJ5zQ4T633347DQ0Nxt/+/fv7cIQSyQBDpvBJJEcF43QBVdlkbhx9Jpx9l7j+5k+gfH0CRibpr5x77rn4/X4CgQBz586NuO2YY47B5XLx8ccfG9sCgQCrV69m4sSJAEyYMKGdCdinn34a8f+xxx7L5s2bGT16dLs/l8vVS89MIjnyJFRA5ebmYrfbqaysjNheWVlJYWFhp/f1er0899xzXH/99Z3u53a7SU9Pj/g7WgmFVdoCfeCyFPQPvCarqhoZXUkwbYEQ4XA/tBU+ylP4Wv3ShUwyOBhbKATU9oqmyBtO/B6MPRdCPnhxAfiaEzA6SX/EbrezZcsWNm/ejN1uj7gtJSWFm266iR/84Ae8+eabbN68mRtvvJGWlhZjHvbd736XHTt28IMf/IBt27bx7LPPsnTp0ojH+dGPfsTKlStZtGgRa9euZceOHbzyyivSREIy4EiogHK5XMycOZPly5cb28LhMMuXL2+XNxvNCy+8gM/niyhyHOx886nPOOU379HU1oviJhyCx2bD46dCONx7xznSvH4b/GYkHNqe6JHQ0BrgxHuXc8NfP0/0UNpzFKfwrT9Qz9S73uL+t7cleigSSa+jR6C2V0UJKEWBeY9B+hCo2Sm+GyUSjc4Wmu+77z4uu+wyvvnNb3Lssceyc+dO3nrrLbKysgAYNmwY//73v1m2bBnTpk3j8ccf59e//nXEY0ydOpUPPviA7du3c+qppzJjxgzuuOMOiouLe/25SSRHkoSaSAAsXryYb33rWxx33HGccMIJPPTQQ3i9XsM+89prr2XIkCHce++9Efd76qmnmDdvHjk5OYkYdr9DVVVW760lEFLZXNbIrFG99Lq01okfXYDWWkjJ7Xz//sLOd8WKa9kayBub0KHsPtRMfUuA1XtrEzqOdoTDEDQbKdJSLbbZBnSmr8GGgw0EQipf7KtL9FAkkl5nbGEqAHurvbQFQniclohCcjZc9hQsPR/W/RNmL4TCKQkaqSSRREeIolm2bJlx3ePx8PDDD/Pwww93uP+FF17IhRdeGLEt2g79+OOP5+233+7wMWL1jpJI+hsJnxldddVV/P73v+eOO+5g+vTprF27ljfffNMwligtLaW8vDziPtu2beOjjz7qMn1vMOELhgmERErYvpqWLvY+DAKWCXbzAEnxCvqhvlRcb6lJ7FiAFi2NrKktSKg/pfEFos6bcBDa6hMylN6gLSAipl6f7IMjOfrJS3WTmewkrMKuQzHS9IbPhonzxPWVj/Tp2CQSiWSgk3ABBbBo0SL27duHz+fjs88+i7DBXLFiRbsVknHjxqGqKmeffXYfj7T/0mhJ29tT04u1PlYBNVBqZOpLQdXSDfuBgLJO4Jvb+tFk3kjfU8CdIa4OFJEcB3p9YLMUUJJBgKIoRj+o7ZVNsXc6+fvicuO/oV4aLEkkkig+/gM8dQ601id6JP2OfiGgJIdPk2Uivq83BZQ1xctb3XvH6YzqHfDxw5FirjNqd5vX+4GAarEYGTS09iMzDt2Bz5kMqfni+lHkxKcLKK+v/xlJVDW18diKXdQ0+7reWSKJk3HRVubRFM+AEaeKaPOnj/XhyCQSyYDg86dh/2ewa3nX+w4ypIA6SrAKqL3VR3kK37u/gHd+Dltfi29/q4BKlOizYI2ANPam4Ud30R34XFYBdfREoHQHvv6Ywrf047385s2t/OWjPYkeiuQownDi6ygCBXDyzeJyzTNylVkikUTSqtUMV2xI7Dj6IVJAHSVYnff21XhR1V6qrekPKXx6PVNTeef76UREoBJv3NDiNyfw/SoCpafwuVIgJU9cP4qc+NqCmoDyB3vv89FDar1+ADYebEjwSCRHE2YEqhMBNfosyJ8I/mb4fEkfjUwikfR7wmFo036TpIBqhxRQRwnWCJTXH+JQb6UCRQioBE2u9eO2xummVrvLvN4PUvisKWSN/UlAGSl8FgF1FKXwtfpFHVxYhda+6JfWDfSo5NbOJrqSPuHee+/l+OOPJy0tjfz8fObNm8e2bZ1b3y9duhRFUSL+PB5PH424Y8YWCCe+g/WtHbe3UBQ4SauF+uxxCMo0UolEAvgaAW2xUQqodkgBdZQQ/ePYa0581hqoREQnVNWc1McbTep3NVD9NQJ1dKfw6REo6H9GEnpd3KEmnxGNkiSGDz74gIULF/Lpp5/yzjvvEAgEOOecc/B6O68tTU9Line truncated
+      "text/plain": [
+       "<Figure size 1000x300 with 2 Axes>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "model_type = \"boomerang\"\n",
+    "task = 1\n",
+    "max_iters = 2000\n",
+    "thr = False\n",
+    "logdir = os.path.join(\"tmp\", f\"task{task}-{max_iters}-{model_type}-{thr}\")\n",
+    "df_mamba = pd.read_pickle(logdir)\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 2, figsize=(10, 3), sharex=True)\n",
+    "sns.lineplot(df_mamba, x=\"Epoch\", y=\"Acc\", hue=\"Model\", ax=ax[0])\n",
+    "sns.lineplot(df_mamba, x=\"Epoch\", y=\"Loss\", hue=\"Model\", ax=ax[1])\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "a222fb0e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "----------  Running in parallel mode  ----------\n",
+      "Accuracy:  0.899\n",
+      "----------  Running in sequential mode  ----------\n",
+      "Accuracy:  0.9\n"
+     ]
+    }
+   ],
+   "source": [
+    "import equinox as eqx\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import jax.random as jrandom\n",
+    "from train import Boomerang, Implicit, ReasoningDataset, SequenceClassifier\n",
+    "\n",
+    "\n",
+    "@eqx.filter_jit\n",
+    "def compute_accuracy(model, inputs, targets, masks):\n",
+    "    def forward_single(inp, tgt, msk):\n",
+    "        logits = model(inp)\n",
+    "        preds = jnp.argmax(logits, axis=-1)\n",
+    "        correct = (preds == tgt) * msk\n",
+    "        return correct.sum(), msk.sum()\n",
+    "\n",
+    "    correct, total = jax.vmap(forward_single)(inputs, targets, masks)\n",
+    "    return correct.sum() / (total.sum() + 1e-8)\n",
+    "\n",
+    "\n",
+    "@eqx.filter_jit\n",
+    "def compute_accuracy_seq(model, inputs, targets, masks):\n",
+    "    def forward_single(inp, tgt, msk):\n",
+    "        logits = model.sequential(inp)\n",
+    "        preds = jnp.argmax(logits, axis=-1)\n",
+    "        correct = (preds == tgt) * msk\n",
+    "        return correct.sum(), msk.sum()\n",
+    "\n",
+    "    correct, total = jax.vmap(forward_single)(inputs, targets, masks)\n",
+    "    return correct.sum() / (total.sum() + 1e-8)\n",
+    "\n",
+    "\n",
+    "key = jrandom.key(0)\n",
+    "keys = jrandom.split(key, 2)\n",
+    "\n",
+    "dataset = ReasoningDataset()\n",
+    "\n",
+    "implicit_model = SequenceClassifier(\n",
+    "    vocab_size=dataset.VOCAB_SIZE,\n",
+    "    d_model=ReasoningDataset.NUM_OUTPUT,\n",
+    "    d_state=16,\n",
+    "    d_inner=128,\n",
+    "    model_class=Implicit if model_type == \"implicit\" else Boomerang,\n",
+    "    max_iters=max_iters,\n",
+    "    # with_thr=thr,\n",
+    "    key=keys[0],\n",
+    ")\n",
+    "\n",
+    "implicit_model = eqx.tree_deserialise_leaves(f\"{logdir}.eqx\", implicit_model)\n",
+    "\n",
+    "eval_inputs, eval_targets, eval_masks = dataset.generate_batch(\n",
+    "    keys[1], 1000, \"simple\" if task == 1 else \"complex\"\n",
+    ")\n",
+    "\n",
+    "print(10 * \"-\", \" Running in parallel mode \", 10 * \"-\")\n",
+    "acc = compute_accuracy(implicit_model, eval_inputs, eval_targets, eval_masks)\n",
+    "print(\"Accuracy: \", acc)\n",
+    "\n",
+    "print(10 * \"-\", \" Running in sequential mode \", 10 * \"-\")\n",
+    "acc = compute_accuracy_seq(implicit_model, eval_inputs, eval_targets, eval_masks)\n",
+    "print(\"Accuracy: \", acc)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "701c2110",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "felice",
+   "language": "python",
+   "name": "python3"
+  },
+  "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.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
@@ -0,0 +1,144 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "dbc7a212",
+   "metadata": {},
+   "source": [
+    "# Simultaneuos & sequential mode\n",
+    "\n",
+    "In this script we test that both simultaneous and sequential model perform the same output with the Boomerang neuron model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "868157a6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax.random as jrandom\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from felice.networks.implicit import Boomerang"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "51adce89",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = jrandom.randint(jrandom.key(0), (20, 1), 0, 15)\n",
+    "model = Boomerang(\n",
+    "    d_model=1,\n",
+    "    d_state=16,\n",
+    "    d_inner=2,\n",
+    "    max_iters=None,\n",
+    "    dt=1e-3,\n",
+    "    key=jrandom.key(42),\n",
+    ")\n",
+    "out_parallel = model(x)\n",
+    "out_sequential = model.sequential(x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "8af5dc49",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import equinox as eqx\n",
+    "import jax.numpy as jnp\n",
+    "\n",
+    "\n",
+    "def loss_fn(model, x):\n",
+    "    output = model(x)\n",
+    "    return jnp.sum(output[-1])\n",
+    "\n",
+    "\n",
+    "gradient = eqx.filter_grad(loss_fn)(model, x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "3ab81e3c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA0wAAAESCAYAAADOhlUZAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAX6lJREFUeJzt3Xd8FHX+x/HX7ia7KaQAIZXQewklaAhgJQiICnYUxV4QPQQLoD9BD4RTznJW1BNErMgpcoIgBAGRKr1DqKEkoaWXze7O7w/u4kUSIEgyKe/n4zEPyOx3vvOeYdjNZ2fmOxbDMAxERERERETkDFazA4iIiIiIiFRWKphERERERERKoYJJRERERESkFCqYRERERERESqGCSUREREREpBQqmEREREREREqhgklERERERKQUXmYHqCgej4cjR44QEBCAxWIxO46IiIiIiJjEMAyysrKIjIzEaj37OaQaUzAdOXKE6Ohos2OIiIiIiEglkZycTP369c/apsYUTAEBAcDpnRIYGGhyGhERERERMUtmZibR0dFFNcLZ1JiC6b+X4QUGBqpgEhERERGR87pVR4M+iIiIiIiIlEIFk4iIiIiISClUMImIiIiIiJRCBZOIiIiIiEgpVDCJiIiIiIiUQgWTiIiIiIhIKVQwiZRRdsoenNmnzI4hIiIiIhVABZNIGRzesRr7+5dy4O1rzY4iIiIiIhVABZNIGexY8g12i4vmBdtI3r7a7DgiIiIiUs5UMImcp5wCF08eSWCrpyEAqb9MMzmRiIiIiJQ3FUwi5+n7DUfIKnDxtvtmABoenYvhdpmcSkRERETKkwomkfNgeNz88OtaAGKuuoUMw596xkl2rfrR5GQiIiIiUp4uqGB69913adSoET4+PsTFxbF6den3clx55ZVYLJYzpn79+hW1MQyDMWPGEBERga+vLwkJCezevbtYPydPnmTQoEEEBgYSHBzMAw88QHZ29oXEFymz3Sv+zacZ9/Oa/UMGdWvB1jo9Ache84XJyURERESkPJW5YPr6668ZMWIEY8eOZd26dXTo0IHevXuTlpZWYvtvv/2Wo0ePFk1btmzBZrNx6623FrV59dVXeeutt5g8eTKrVq3C39+f3r17k5+fX9Rm0KBBbN26lQULFvDDDz+wdOlSHn744QvYZJGyy1/xIV4WD2EhdQny86ZWl0GkGLVZdaoWTpfH7HgiIiIiUk4shmEYZVkgLi6OSy65hHfeeQcAj8dDdHQ0TzzxBKNGjTrn8m+++SZjxozh6NGj+Pv7YxgGkZGRPPXUUzz99NMAZGRkEBYWxieffMLAgQPZvn07bdq0Yc2aNXTp0gWAefPmce2113Lo0CEiIyPPud7MzEyCgoLIyMggMDCwLJssNdzJI3sI+iAWm8Vg9y2JNG/XBbfbQ4+/LeBolosP7o6ld9tws2OKiIiIyHkqS21QpjNMTqeTtWvXkpCQ8HsHVisJCQmsWLHivPr4+OOPGThwIP7+/gDs27ePlJSUYn0GBQURFxdX1OeKFSsIDg4uKpYAEhISsFqtrFq1qsT1FBQUkJmZWWwSuRB7572LzWKwyTuG5u1OH4M2m5XrOzUAYNb6w2bGExEREZFyVKaC6fjx47jdbsLCworNDwsLIyUl5ZzLr169mi1btvDggw8WzfvvcmfrMyUlhdDQ0GKve3l5UadOnVLXO3HiRIKCgoqm6Ojoc2+gyB+4CwtocnAmADkx9xV7bUDHKGy4yd+xgMyTqWbEExEREZFyVqGj5H388ce0b9+eSy+9tNzXNXr0aDIyMoqm5OTkcl+nVD/bEj+jDhmkUZtOve4s9lqbyEC+qvUGU70msnvhVJMSioiIiEh5KlPBFBISgs1mIzW1+LfpqamphIef/R6OnJwcvvrqKx544IFi8/+73Nn6DA8PP2NQCZfLxcmTJ0tdr8PhIDAwsNgkUlaWDZ8BsDPqZnx8fM543dnoagACdn9XoblEREREpGKUqWCy2+3ExsaSmJhYNM/j8ZCYmEh8fPxZl/3mm28oKCjgrrvuKja/cePGhIeHF+szMzOTVatWFfUZHx9Peno6a9euLWqzaNEiPB4PcXFxZdkEkfO2/3gOd6Q/xouF99D4msdKbNPs6ntwGVZaFO4gZe+WCk4oIiIiIuWtzJfkjRgxgo8++ohp06axfft2hgwZQk5ODvfdd/r+jsGDBzN69Ogzlvv4448ZMGAAdevWLTbfYrHw5JNPMn78eGbPns3mzZsZPHgwkZGRDBgwAIDWrVvTp08fHnroIVavXs2vv/7K448/zsCBA89rhDyRC/H5qgNk4cf+ZndRv2HTEtuERTZgi09nAA4umVaR8URERESkAniVdYHbb7+dY8eOMWbMGFJSUujYsSPz5s0rGrTh4MGDWK3F67CdO3eybNkyfvrppxL7fPbZZ8nJyeHhhx8mPT2dHj16MG/evGKXQH3++ec8/vjj9OzZE6vVys0338xbb71V1vgi5yXfWciM3w4BcHfXhmdv2/oW2PAbUQdnY3hewWKt0FsDRURERKQclfk5TFWVnsMkZbHum1ewbf6KL31u4+VRo7BZLaW2zcxMx/ZaC/wtBey5YRZNO19VgUlFREREpKzK7TlMIjWCYVBvx3Q6WPfSr4HnrMUSQGBgMFsDLwfg0Jp/V0RCEREREakgZb4kT6S6S1ozj2buZHIMB237Pnxey3i6D+e6768gJa0VK90evGz6LkJERESkOtBvdSJ/kLPsAwA21u5Nnbr1zmuZ2EviOezbkuM5Tn7dc6I844mIiIhIBVLBJPI/MlIP0iZjKQDBVww57+W8bVau73B6xMbZ6w6USzYRERERqXgqmET+x+557+JtcbPFqw2tO5792WJ/dFP7Okzymszo7TeRk3GynBKKiIiISEVSwSTyH55CJw33zwAgs909WCxnH+zhjzo0CucS772EWDLY8fPn5RFRRERERCqYCiaR//hlz0meK7iX+XSl4zV3l3l5i9XK4QY3AOCzfebFjiciIiIiJlDBJPIf01cdYoGnCytjX8fPz/+C+mh45T0AtM7fyPEjey9mPBERERExgQomEeDQqVwW7UgF4K6uDS+4n/qNW7HVuy1Wi8GeRdMuVjwRERERMYkKJhFg37cv8hfbTPo1hqb1av2pvrKa3wRA6L5ZFyGZiIiIiJhJBZPUeAW5GXRKns6TXt9yb9PcP91fy56DcRpeNHbv58C2VRchoYiIiIiYxcvsACJm2/7TFDqSSzLhdLryxj/dX+26oSwI6s/mExYcSYUMbXMRQoqIiIiIKXSGSWo2wyB4y+l7jfY1HoiX18X5DsGZ8DJvuW/ii21OPB7jovQpIiIiIhVPBZPUaPvW/0wj1z7yDW9a93n0ovXbs3UoAQ4vDqfnsWa/HmIrIiIiUlWpYJIaLX3p+wCsD+pJvbCIi9avj7eNG9rWoY91NSkL375o/YqIiIhIxdI9TFJjZZ04QttTi8ACtS4bctH7v6P+cdpte5PsI77k5z6Nj9+fG31PRERERCqezjBJjTVvYzIz3Zex3taedl2uuOj9t4m7hqPUoxZ5bF8y46L3LyIiIiLlTwWT1EiGYfDBhnyecz3E5p7TsVgsF30dVpuNfZH9Tv99swomERERkapIBZPUSCv3niQpLRs/u40bO9cvt/VEXjYYgDY5q0k/drTc1iMiIiIi5UMFk9RIh3+cRIxlDzd2iiLAx7vc1tOodSy7bU3xtrjZ9fOn5bYeERERESkfKpikxjl+cAe3HH+fWfYx3Nu2/Mc9Od54AABBu78r93WJiIiIyMWlgklqnAPz3wJggyOW5i1al/v6ml19D27DQl5BAckpx8t9fSIiIiJy8ahgkhqlMD+HZoe/P/33zvdXyDrrRTZkeORnDHCO57utpypknSIiIiJycahgkhpl+4JPCCKbI9Sj09W3Vdh6r+jSEYBZ6w9jGEaFrVdERERE/hwVTFKj+G/6BICkBrdht5ffYA9/1KddOL7eNo4fT2Pbrt0Vtl4RERER+XNUMEmNkbz5F5oW7qLA8KJ5nyEVum5/hxfjIn5ljeMxshe+UqHrFhEREZELp4JJaoxfNmzniFGHdQFXEhEZXeHrb966Aw5LIc2P/UShs6DC1y8iIiIiZaeCSWqEnAIXE5MaclnBPzB6/82UDG173MAJgqhDJtuWfW9KBhEREREpGxVMUiN8v+EIWQUuousG0LVtM1MyeHnb2R3aG4DC9V+ZkkFEREREykYFkwk8HoN/rT3EiWxdllURDI+bnUtn4IWLu7o2xGq1mJalbvzdALTN/IWsjJOm5RARERGR86OCyQTPz9rMlJnfs/KzF82OUiMkrZzNSznj+NHxHLd2jjI1S7MOPThojcLX4mT7z1+amkVEREREzu2CCqZ3332XRo0a4ePjQ1xcHKtXrz5r+/T0dIYOHUpERAQOh4MWLVowd+7cotezsrJ48sknadiwIb6+vnTr1o01a9YU6+Pee+/FYrEUm/r06XMh8U13R0sL39v/j34p77Fr7c9mx6n28pd/BEBKSFeC/B2mZrFYrRxucAMAvttnmppFRERERM6tzAXT119/zYgRIxg7dizr1q2jQ4cO9O7dm7S0tBLbO51OevXqxf79+5k5cyY7d+7ko48+Iirq92/6H3zwQRYsWMD06dPZvHkz11xzDQkJCRw+fLhYX3369OHo0aNF05dfVs1v6GPatmdd7dPFnvXHp/G4XCYnqr5OHk6iTdZyAEKvGmpymtMaXnkf4wrv4oHMB0jJyDc7joiIiIicRZkLptdff52HHnqI++67jzZt2jB58mT8/PyYMmVKie2nTJnCyZMnmTVrFt27d6dRo0ZcccUVdOjQAYC8vDz+9a9/8eqrr3L55ZfTrFkzXnzxRZo1a8b7779frC+Hw0F4eHjRVLt27QvY5Mqh8cBXyDT8aOZK4rdZb5kdp9raN/8dbBaDjd4dadmus9lxAIhs1JJLine truncated
+      "text/plain": [
+       "<Figure size 1000x300 with 1 Axes>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "t = -1\n",
+    "fig, ax = plt.subplots(1, 1, figsize=(10, 3))\n",
+    "ax.plot(out_parallel)\n",
+    "ax.plot(out_sequential, \"--\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "0c2e88d4",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiwAAAGdCAYAAAAxCSikAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAPYQAAD2EBqD+naQAARMBJREFUeJzt3Xl8FfW9//HXnD37SlZCwiqI7IGIG1ZTqNrVapVSRe3tal2KbcXbW7D1Z0HrbV1QtGrVVr0urdwqvSKIgFVRNpF9FQJkD5CT5CQ56/z+iMZSQLKcQ4bwfj4e5wE5Z+b7/X4IMO/MfOc7hmmaJiIiIiIWZuvpAYiIiIiciAKLiIiIWJ4Ci4iIiFieAouIiIhYngKLiIiIWJ4Ci4iIiFieAouIiIhYngKLiIiIWJ6jpwcQLZFIhIqKCpKSkjAMo6eHIyIiIh1gmiaNjY3k5eVhsx3/PEqvCSwVFRUUFBT09DBERESkC/bv30/fvn2P+3mvCSxJSUlAW8HJyck9PBoRERHpiIaGBgoKCtqP48fTawLLp5eBkpOTFVhEREROMSeazqFJtyIiImJ5CiwiIiJieQosIiIiYnkKLCIiImJ5CiwiIiJieQosIiIiYnkKLCIiImJ5CiwiIiJieQosIiIiYnkKLCIiImJ5CiwiIiJieQosIiIiYnm95uGHIiIivUm4qYnmVavwvfsecWPHkHLZZT09pB6lwCIiImIRgf37aVi4kKa3/0nLhg0QDgMQrKxsDyymaRI8cIBA2T6CB/YT2L+fUE0tkcZGIk1NhJuaIBwCwwZ2O4ZhYEtMxJ6ejj0tFUdaOs6CAtwDB+Dq3x97UlJPltxhCiwiIiI9yDRNfO+8y8E/PUnzyveP+MxZ2I+Ec87BmZtH3fz5NK9fT+v6jwh7vVHr35GTQ/y4ccSXTCChpARXYWHU2o4mwzRNs6cHEQ0NDQ2kpKTg9XpJTk7u6eGIiIicUGDfPipnz/4sqNhsJEycSNKUycSNHEnTsmV4//4qgb17T9qY3MOGkXr55SR/+TIcaWkx76+jx28FFhERkR7ge+89Dtx0MxGfD8PlIm3q1aRPn44jO5tDTz9D3WOPEWloAMBwubBnZhCqqIzZeJz9+hGsrIRgsL3PtGnTyLzxx9gTE2PWrwKLiIiIRfl37WLPFVditrYSN24ceXPn4CoowAwEOHDLrTQtWwaAa9BA0q+9lpa1a/H+/dXYDsrpJHnyZDzDhuL9x//h37oVAEduLn3nPUTc8OEx6bajx2/d1iwiInKS1dx/P2ZrK/ETz6bfU3/CVVAAQN0TT9C0bBmG203OXb9hwN//TvOq1e1hJW7MmJiMx1lQAMEgDf/4B96//52Cxx6l4LFHcRYUEKqsZN/06/B/vCcmfXeUAouIiMhJ1rxqNQBZM27D5nK1v9/w6msA5MyeTdqVV+LfvZuGhQvBZqPvo/Nx5ubEZDz29DQKnnwCe59M/Dt3UTX7ThInTaL/gleIGzOGSFMT1b/9bUz67igFFhERkZPMcLeFlH+/28eMRI74PNLYCIAtMRHPsDMJlJdHdRx9brkZgNaPNuDMzqbfE08C0LRiBRG/H3tiInlz5wDge/ddIq2tUe2/MxRYRERETrKk0lIAqu+6i1BdXfv7yZdeAkDVnb+m8a1lxI0YgauwkEhDA3u/9S0MI7qH7ebVa4BPJvVmZBDYs6f9a8Nub9vI/skKKKbZ9uohCiwiIiInWZ+bbsKRl0ugrIw9V1xJ07vvApD5wx8SX1xMpLGRAz/+Mftv/AkZP/ohrsJCQtXVtKxfH9Vx+N57D4D0G67n0DPPUD5jBgBpV12F4XAQbvJROXMmAPElJdji4qLaf2foLiEREZEeECgrY/8Pf9R+ViNh0gVkXHcdcWPGUPvgQxz6y1/abzF2ZGcTqq6OyThsSUlgmkSamgBI+ebl5M6eje+DVVT9v7sIlu3DFh9P0Ysv4B48OOr967ZmERERiws3+ah94AEOP/98+zL8zvx8Ui7/Bu7Bg2lasYLG1xcRaW6O+VjcgweTNu3bGE4Xh196kdaPNgBtYanvvIeIGzEiJv0qsIiIiJwiAmVlHHzqKRoW/qP9TAeAPSODuFGjCDd4CdfWETp48IjPu8uIi8OZk4MjJxszEKTlww/h04m/Tidp355K5o03Yo/hcVWBRURE5BQTaWmhcckSvAsX0rx6DWZLS+w6s9naw8m/cg8dSvKXppB65ZU4MjJi1/8nFFhEREROYWYgQMuGDfjee4+WDRvx79hBqKYm6v04CwrwnHkm8ePGkXjRRbj65ke9j8/T0eO3ntYsIiJiQYbLRXxxMfHFxe3vhevr8e/cSaCsjFBtLaHaWoI1NYQPHiLib8X0BzD9fiL+Vgy7A5vbjREXh83jwZaYiDM3B0dODs6cHJz5ffEMG4o9JaUHq+w4BRYREZFThD01lfjx44kfP76nh3LSaR0WERERsTwFFhEREbE8BRYRERGxPAUWERERsTwFFhEREbE8BRYRERGxPAUWERERsTwFFhEREbE8BRYRERGxvC4FlocffpiioiI8Hg8lJSWsWrXquNs+/vjjnH/++aSlpZGWlkZpaennbv/DH/4QwzC4//77uzI0ERER6YU6HVhefPFFZsyYwezZs1m3bh2jRo1iypQp1BzngUzLly9n6tSpLFu2jJUrV1JQUMDkyZMpLy8/atsFCxbw/vvvk5eX1/lKREREpNfqdGD5/e9/z/e+9z2uv/56zjzzTB599FHi4+P505/+dMztn3vuOX784x8zevRohg4dyhNPPEEkEmHp0qVHbFdeXs5NN93Ec889h9Pp7Fo1IiIi0it1KrAEAgHWrl1LaWnpZw3YbJSWlrJy5coOtdHc3EwwGCQ9Pb39vUgkwjXXXMPPf/5zhg8f3qF2/H4/DQ0NR7xERESkd+pUYKmrqyMcDpOdnX3E+9nZ2VRVVXWojdtvv528vLwjQs8999yDw+Hg5ptv7vBY5syZQ0pKSvuroKCgw/uKiIjIqeWk3iU0d+5cXnjhBRYsWIDH4wFg7dq1PPDAAzz99NMYhtHhtu644w68Xm/7a//+/bEatoiIiPSwTgWWzMxM7HY71dXVR7xfXV1NTk7O5+573333MXfuXBYvXszIkSPb3//nP/9JTU0N/fr1w+Fw4HA4KCsr47bbbqOoqOi47bndbpKTk494iYiISO/UqcDicrkYN27cERNmP51AO3HixOPud++993LXXXexaNEiiouLj/jsmmuuYcOGDaxfv779lZeXx89//nPeeOONTpYjIiIivZGjszvMmDGD6dOnU1xczIQJE7j//vvx+Xxcf/31AFx77bXk5+czZ84coG1+yqxZs3j++ecpKipqn+uSmJhIYmIiGRkZZGRkHNGH0+kkJyeHM844o7v1iYiISC/Q6cBy1VVXUVtby6xZs6iqqmL06NEsWrSofSLuvn37sNk+O3Ezf/58AoEAV1xxxRHtzJ49mzvvvLN7oxcREZHTgmGaptnTg4iGhoYGUlJS8Hq9ms8iIiJyiujo8VvPEhIRERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHLU2ARERERy1NgEREREctTYBERERHL61JgefjhhykqKsLj8VBSUsKqVauOu+3jjz/O+eefT1paGmlpaZSWlh6xfTAY5Pbbb2fEiBEkJCSQl5fHtddeS0VFRVeGJiIiIr1QpwPLiy++yIwZM5g9ezbr1q1j1KhRTJkyhZqammNuv3z5cqZOncqyZctYuXIlBQUFTJ48mfLycgCam5tZt24dv/rVr1i3bh2vvPIK27dv56tf/Wr3KhMREZFewzBN0+zMDiUlJYwfP5558+YBEIlEKCgo4KabbmLmzJkn3D8cDpOWlsa8efO49tprj7nN6tWrmTBhAmVlZfTr169D42poaCAlJQWv10tycnLHCxIREZEe09Hjd6fOsAQCAdauXUtpaelnDdhslJaWsnLlyg610dzcTDAYJD09/bjbeL1eDMMgNTW1M8MTERGRXsrRmY3r6uoIh8NkZ2cf8X52djbbtm3rUBu33347eXl5R4Sef9Xa2srtt9/O1KlTPzdp+f1+/H5/+9cNDQ0d6l9EREROPSf1LqG5c+fywgsvsGDBAjwez1GfB4NBvvWtb2GaJvPnz//ctubMmUNKSkr7q6CgIFbDFhERkR7WqcCSmZmJ3W6nurr6iPerq6vJycn53H3vu+8+5s6dy+LFixk5cuRRn38aVsrKyliyZMkJ56HccccdeL3e9tf+/fs7U4qIiIicQjoVWFwuF+PGjWPp0qXt70UiEZYuXcrEiROPu9+9997LXXfdxaJFiyguLj7q80/Dys6dO3nzzTfJyMg44VjcbjfJyclHvERERKR36tQcFoALine truncated
+      "text/plain": [
+       "<Figure size 640x480 with 1 Axes>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "ts, ys = model.get_z(x, t0=2, t1=4, num_points=1000)\n",
+    "plt.figure()\n",
+    "plt.plot(ys[:, :, 0], ys[:, :, 1])\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "79112537",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "felice",
+   "language": "python",
+   "name": "python3"
+  },
+  "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.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
@@ -0,0 +1,281 @@
+import argparse
+import os
+from typing import Any, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+import optax
+import pandas as pd
+from jaxtyping import Array, Float
+from optax import OptState
+from tqdm import trange
+
+from felice.datasets.reasoning import ReasoningDataset
+from felice.networks import Implicit, Mamba, SequenceClassifier
+from felice.networks.implicit.boomerang import ImplicitBoomerang
+
+
+def compute_loss(
+    model: eqx.Module, inputs: Array, targets: Array, masks: Array
+) -> Float[Array, ""]:
+    def forward_single(inp, tgt, msk):
+        logits = model(inp)
+        loss = optax.softmax_cross_entropy_with_integer_labels(logits, tgt)
+        return (loss * msk).sum() / (msk.sum() + 1e-8)
+
+    losses = jax.vmap(forward_single)(inputs, targets, masks)
+    return losses.mean()
+
+
+v_and_grad = eqx.filter_value_and_grad(compute_loss)
+
+
+@eqx.filter_jit
+def compute_accuracy(
+    model: eqx.Module, inputs: Array, targets: Array, masks: Array
+) -> Float[Array, ""]:
+    def forward_single(inp, tgt, msk):
+        logits = model(inp)
+        preds = jnp.argmax(logits, axis=-1)
+        correct = (preds == tgt) * msk
+        return correct.sum(), msk.sum()
+
+    correct, total = jax.vmap(forward_single)(inputs, targets, masks)
+    return correct.sum() / (total.sum() + 1e-8)
+
+
+@eqx.filter_jit
+def train_step(
+    model: eqx.Module,
+    opt_state: OptState,
+    optimizer: Any,
+    inputs: Array,
+    targets: Array,
+    masks: Array,
+) -> Tuple[eqx.Module, OptState, Array]:
+    loss, grads = v_and_grad(model, inputs, targets, masks)
+    updates, opt_state = optimizer.update(grads, opt_state, model)
+    model = eqx.apply_updates(model, updates)
+    return model, opt_state, loss
+
+
+def train_and_compare(
+    model_type: Any,
+    logdir: str,
+    task_type: str = "simple",
+    n_epochs: int = 1000,
+    batch_size: int = 64,
+    d_model: int = 64,
+    d_state: int = 16,
+    d_inner: int = 32,
+    dt: float = 1.0,
+    max_iters: int = 8,
+    lr: float = 1e-3,
+    seed: int = 42,
+    # with_thr: bool = True,
+) -> Tuple[eqx.Module, eqx.Module, Array, Array, pd.DataFrame]:
+    r"""Train Mamba and implicit model on the reasoning synthetic dataset.
+
+    Args:
+        model_type: The type of the implicit model to train (Boomerang, Mamba Implicit).
+        logdir: Directory and filenmae of the log.
+        task_type: Type of task to solve from the reasoning synthetic dataset (simple, accumulation).
+        n_epochs: Number of epochs to train.
+        batch_size: Training batch size.
+        d_model: Model dimensions including output.
+        d_state: Model state dimension.
+        d_inner: Model latent dimension.
+        max_iters: Maximum number of iterations in the implicit model.
+        lr: Learning rate.
+        seed: Random seed.
+        with_thr: For the Boomerang model, if using threshold for dual fixpoints.
+
+    Returns:
+        The trained models (mamba and implicit) with the respective final accuracy and
+        a pandas dataframe with the loss and accuracy per epoch.
+    """
+    key = jrandom.key(seed)
+    keys = jrandom.split(key, 4)
+
+    dataset = ReasoningDataset()
+
+    standard_model = SequenceClassifier(
+        vocab_size=dataset.VOCAB_SIZE,
+        d_model=d_model,
+        d_state=d_state,
+        d_inner=d_inner,
+        model_class=Mamba,
+        key=keys[0],
+    )
+
+    implicit_model = SequenceClassifier(
+        vocab_size=dataset.VOCAB_SIZE,
+        d_model=d_model,
+        d_state=d_state,
+        d_inner=d_inner,
+        model_class=model_type,
+        max_iters=max_iters,
+        dt=dt,
+        # with_thr=with_thr,
+        key=keys[1],
+    )
+
+    # implicit_model = ImplicitBoomerang(
+    #     vocab_size=dataset.VOCAB_SIZE,
+    #     d_model=d_model,
+    #     d_state=d_state,
+    #     d_inner=d_inner,
+    #     max_iters=max_iters,
+    #     dt=dt,
+    #     # with_thr=with_thr,
+    #     key=keys[1],
+    # )
+    # Count parameters
+    def count_params(model):
+        return sum(
+            x.size for x in jax.tree_util.tree_leaves(eqx.filter(model, eqx.is_array))
+        )
+
+    print(f"Mamba SSM params: {count_params(standard_model):,}")
+    print(f"Implicit SSM params: {count_params(implicit_model):,}")
+
+    # Optimizers
+    optimizer = optax.adam(lr)
+    standard_opt_state = optimizer.init(eqx.filter(standard_model, eqx.is_array))
+    implicit_opt_state = optimizer.init(eqx.filter(implicit_model, eqx.is_array))
+
+    # Training loop
+    print(f"\nTraining on task: {task_type} with {max_iters} steps")
+    print("=" * 60)
+
+    train_key = keys[2]
+
+    df = pd.DataFrame({"Epoch": [], "Loss": [], "Acc": [], "Model": []})
+    pbar = trange(n_epochs)
+    for epoch in pbar:
+        train_key, batch_key = jrandom.split(train_key)
+        inputs, targets, masks = dataset.generate_batch(
+            batch_key, batch_size, task_type
+        )
+
+        # Train standard model
+        standard_model, standard_opt_state, standard_loss = train_step(
+            standard_model, standard_opt_state, optimizer, inputs, targets, masks
+        )
+
+        # Train implicit model
+        implicit_model, implicit_opt_state, implicit_loss = train_step(
+            implicit_model, implicit_opt_state, optimizer, inputs, targets, masks
+        )
+
+        if (epoch + 1) % 10 == 0:
+            # Evaluate on fresh batch
+            eval_key = jrandom.fold_in(keys[3], epoch)
+            eval_inputs, eval_targets, eval_masks = dataset.generate_batch(
+                eval_key, batch_size, task_type
+            )
+
+            standard_acc = compute_accuracy(
+                standard_model, eval_inputs, eval_targets, eval_masks
+            )
+            implicit_acc = compute_accuracy(
+                implicit_model, eval_inputs, eval_targets, eval_masks
+            )
+
+            new_df = pd.DataFrame(
+                {
+                    "Epoch": [epoch, epoch],
+                    "Loss": [standard_loss.item(), implicit_loss.item()],
+                    "Acc": [standard_acc.item(), implicit_acc.item()],
+                    "Model": ["Mamba", "Implicit"],
+                }
+            )
+            df = pd.concat([df, new_df], ignore_index=True)
+            df.to_pickle(logdir)
+            pbar.write(
+                f"Epoch {epoch + 1:4d} | "
+                f"Standard: loss={standard_loss:.4f}, acc={standard_acc:.4f} | "
+                f"Implicit: loss={implicit_loss:.4f}, acc={implicit_acc:.4f}"
+            )
+
+    # Final evaluation
+    print("\n" + "=" * 60)
+    print("Final Evaluation (1000 samples)")
+    print("=" * 60)
+
+    eval_inputs, eval_targets, eval_masks = dataset.generate_batch(
+        keys[4], 1000, task_type
+    )
+
+    standard_acc = compute_accuracy(
+        standard_model, eval_inputs, eval_targets, eval_masks
+    )
+    implicit_acc = compute_accuracy(
+        implicit_model, eval_inputs, eval_targets, eval_masks
+    )
+
+    print(f"Mamba SSM accuracy: {standard_acc:.4f}")
+    print(f"Implicit SSM accuracy: {implicit_acc:.4f}")
+    print(f"Improvement: {(implicit_acc - standard_acc) * 100:.2f}%")
+
+    return standard_model, implicit_model, standard_acc, implicit_acc, df
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-t", type=int, choices=[1, 2], default=1, help="Task to perform"
+    )
+    parser.add_argument(
+        "-m",
+        type=str,
+        choices=["boomerang", "implicit"],
+        default="implicit",
+        help="Neuron model to use",
+    )
+    parser.add_argument("--dt", type=float, default=0.001, help="Simulation timestep")
+    parser.add_argument("-i", type=int, default=8, help="Maximum number of iterations")
+    parser.add_argument("-b", type=int, default=64, help="Batch size")
+    parser.add_argument(
+        "--thr", action="store_true", help="Using threshold on the boomerang neuron"
+    )
+    args = parser.parse_args()
+
+    if args.m == "boomerang":
+        model_type = ImplicitBoomerang
+    elif args.m == "implicit":
+        model_type = Implicit
+    else:
+        raise NotImplementedError(f"{args.t} model type not implemented")
+
+    logdir = os.path.join("tmp", f"task{args.t}-{args.i}-{args.m}-{args.thr}")
+    if not os.path.exists("tmp"):
+        os.makedirs("tmp")
+
+    print(f"Saving at {logdir}")
+    _, implicit_model, std_acc1, imp_acc1, df = train_and_compare(
+        model_type,
+        logdir,
+        task_type="simple" if args.t == 1 else "accumulation",
+        n_epochs=1000,
+        batch_size=64,
+        d_model=ReasoningDataset.NUM_OUTPUT,
+        d_state=16,
+        d_inner=128,
+        dt=args.dt,
+        max_iters=args.i,
+        # with_thr=args.thr,
+    )
+    eqx.tree_serialise_leaves(f"{logdir}.eqx", implicit_model)
+    df.to_pickle(logdir)
+
+    print("=" * 70)
+    print("SUMMARY")
+    print("=" * 70)
+    print(f"{'Task':<25} {'Mamba SSM':<15} {'Implicit SSM':<15} {'Delta':<10}")
+    print("-" * 70)
+    print(
+        f"{'Simple Comparison':<25} {std_acc1:<15.4f} {imp_acc1:<15.4f} {(imp_acc1 - std_acc1) * 100:>+.2f}%"
+    )