Try a scalar/array variants of the pure numba impl

Author: ricardoV94

doc/gallery/scan/numba_fib_scan.ipynb

@@ -38,126 +38,164 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:20:23.571190Z",
-     "start_time": "2025-10-07T10:20:23.567707Z"
+     "end_time": "2025-10-07T10:32:18.289971Z",
+     "start_time": "2025-10-07T10:32:18.284515Z"
     }
    },
    "cell_type": "code",
    "source": [
     "import numba\n",
     "\n",
-    "@numba.jit(nopython=True)\n",
-    "def fibonacci_numba(b):\n",
+    "@numba.njit\n",
+    "def fibonacci_numba_scalar(b):\n",
     "    b = b.copy()\n",
     "    a = np.ones((), dtype=np.int32)\n",
     "    for _ in range(N_STEPS):\n",
     "        a[()], b[()] = a[()] + b[()], a[()]\n",
+    "    return a\n",
+    "\n",
+    "@numba.njit\n",
+    "def fibonacci_numba_array(b):\n",
+    "    a = np.ones((), dtype=np.int32)\n",
+    "    for _ in range(N_STEPS):\n",
+    "        a, b = np.asarray(a + b), a\n",
     "    return a"
    ],
    "id": "b1d657d366647ada",
    "outputs": [],
-   "execution_count": 2
+   "execution_count": 66
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:20:26.947566Z",
-     "start_time": "2025-10-07T10:20:23.615573Z"
+     "end_time": "2025-10-07T10:32:19.607842Z",
+     "start_time": "2025-10-07T10:32:19.423324Z"
     }
    },
    "cell_type": "code",
    "source": [
     "b = np.ones((), dtype=np.int32)\n",
-    "assert fibonacci_pytensor(b) == fibonacci_numba(b)\n",
-    "assert fibonacci_pytensor_numba(b) == fibonacci_numba(b)"
+    "assert fibonacci_numba_array(b) == fibonacci_numba_scalar(b)"
    ],
    "id": "7f45c87d259852e6",
    "outputs": [],
-   "execution_count": 3
+   "execution_count": 67
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:20:28.819553Z",
-     "start_time": "2025-10-07T10:20:27.015141Z"
+     "end_time": "2025-10-07T10:32:22.705191Z",
+     "start_time": "2025-10-07T10:32:20.090353Z"
     }
    },
    "cell_type": "code",
-   "source": "%timeit fibonacci_pytensor(b)",
-   "id": "90f27f42b275e9da",
+   "source": "%timeit fibonacci_numba_scalar(b)",
+   "id": "b01c8978960c6e3d",
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "2.22 ms ± 37.8 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
+      "3.21 μs ± 20.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
      ]
     }
    ],
-   "execution_count": 4
+   "execution_count": 68
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:20:42.302214Z",
-     "start_time": "2025-10-07T10:20:28.871240Z"
+     "end_time": "2025-10-07T10:32:25.876514Z",
+     "start_time": "2025-10-07T10:32:23.122275Z"
     }
    },
    "cell_type": "code",
-   "source": "%timeit fibonacci_pytensor_numba(b)",
-   "id": "d918043168a39d59",
+   "source": "%timeit fibonacci_numba_array(b)",
+   "id": "bfc8794b219db03e",
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "165 μs ± 468 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
+      "32.8 μs ± 2.48 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
      ]
     }
    ],
-   "execution_count": 5
+   "execution_count": 69
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": [
+    "assert fibonacci_pytensor(b) == fibonacci_numba_scalar(b)\n",
+    "assert fibonacci_pytensor_numba(b) == fibonacci_numba_scalar(b)"
+   ],
+   "id": "a2185c1de1297a11"
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:20:55.256172Z",
-     "start_time": "2025-10-07T10:20:42.355007Z"
+     "end_time": "2025-10-07T10:29:44.724064Z",
+     "start_time": "2025-10-07T10:29:42.655693Z"
     }
    },
    "cell_type": "code",
-   "source": "%timeit fibonacci_pytensor_numba.vm.jit_fn(b)",
-   "id": "69ec1c01dad9fb66",
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "159 μs ± 1.41 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
+      "2.49 ms ± 327 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
      ]
     }
    ],
-   "execution_count": 6
+   "execution_count": 54,
+   "source": "%timeit fibonacci_pytensor(b)",
+   "id": "f1e8bb6a0c673c8f"
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:20:57.954364Z",
-     "start_time": "2025-10-07T10:20:55.346865Z"
+     "end_time": "2025-10-07T10:29:58.922566Z",
+     "start_time": "2025-10-07T10:29:44.752331Z"
     }
    },
    "cell_type": "code",
-   "source": "%timeit fibonacci_numba(b)",
-   "id": "b01c8978960c6e3d",
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "3.2 μs ± 19.2 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
+      "175 μs ± 6.13 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
      ]
     }
    ],
-   "execution_count": 7
+   "execution_count": 55,
+   "source": "%timeit fibonacci_pytensor_numba(b)",
+   "id": "17cd2859b4c6d3bd"
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-10-07T10:30:11.832294Z",
+     "start_time": "2025-10-07T10:29:59.016709Z"
+    }
+   },
+   "cell_type": "code",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "158 μs ± 706 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
+     ]
+    }
+   ],
+   "execution_count": 56,
+   "source": "%timeit fibonacci_pytensor_numba.vm.jit_fn(b)",
+   "id": "6deb056f63953a42"
   },
   {
    "metadata": {
@@ -362,8 +400,8 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:20:58.266498Z",
-     "start_time": "2025-10-07T10:20:58.254087Z"
+     "end_time": "2025-10-07T10:30:25.495418Z",
+     "start_time": "2025-10-07T10:30:25.481386Z"
     }
    },
    "cell_type": "code",
@@ -458,29 +496,29 @@
    ],
    "id": "bcefae049d4d2540",
    "outputs": [],
-   "execution_count": 13
+   "execution_count": 59
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:21:01.526942Z",
-     "start_time": "2025-10-07T10:20:58.310079Z"
+     "end_time": "2025-10-07T10:30:31.559493Z",
+     "start_time": "2025-10-07T10:30:30.263832Z"
     }
    },
    "cell_type": "code",
    "source": [
     "b = np.ones((), dtype=np.int32)\n",
-    "assert comparable_fibonacci_numba(b) == fibonacci_numba(b)"
+    "assert comparable_fibonacci_numba(b) == fibonacci_numba_scalar(b)"
    ],
    "id": "65887ebba21f46c3",
    "outputs": [],
-   "execution_count": 14
+   "execution_count": 60
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-10-07T10:21:06.041171Z",
-     "start_time": "2025-10-07T10:21:01.578757Z"
+     "end_time": "2025-10-07T10:30:35.999409Z",
+     "start_time": "2025-10-07T10:30:31.567997Z"
     }
    },
    "cell_type": "code",
@@ -491,11 +529,11 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "55 μs ± 756 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
+      "54.6 μs ± 1.28 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
      ]
     }
    ],
-   "execution_count": 15
+   "execution_count": 61
   },
   {
    "metadata": {