Skip to content

Commit 08f1777

Browse files
ricardoV94ricardoV94
committed
Add global N_STEPS
1 parent 2b5ec3c commit 08f1777

File tree

1 file changed

+45
-46
lines changed

1 file changed

+45
-46
lines changed

‎doc/gallery/scan/numba_fib_scan.ipynb‎

Lines changed: 45 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,17 @@
66
"metadata": {
77
"collapsed": true,
88
"ExecuteTime": {
9-
"end_time": "2025-10-07T09:52:07.812763Z",
10-
"start_time": "2025-10-07T09:52:05.905923Z"
9+
"end_time": "2025-10-07T10:20:23.561430Z",
10+
"start_time": "2025-10-07T10:20:21.620124Z"
1111
}
1212
},
1313
"source": [
1414
"import pytensor\n",
1515
"import pytensor.tensor as pt\n",
1616
"import numpy as np\n",
1717
"\n",
18+
"N_STEPS = 1000\n",
19+
"\n",
1820
"b_symbolic = pt.scalar(\"b_symbolic\", dtype=\"int32\")\n",
1921
"\n",
2022
"def step(a, b):\n",
@@ -23,7 +25,7 @@
2325
"(outputs_a, outputs_b), _ = pytensor.scan(\n",
2426
" fn=step,\n",
2527
" outputs_info=[pt.constant(1, dtype=\"int32\"), b_symbolic],\n",
26-
" n_steps=10\n",
28+
" n_steps=N_STEPS\n",
2729
")\n",
2830
"\n",
2931
"# compile function returning final a\n",
@@ -36,8 +38,8 @@
3638
{
3739
"metadata": {
3840
"ExecuteTime": {
39-
"end_time": "2025-10-07T09:52:07.821649Z",
40-
"start_time": "2025-10-07T09:52:07.817580Z"
41+
"end_time": "2025-10-07T10:20:23.571190Z",
42+
"start_time": "2025-10-07T10:20:23.567707Z"
4143
}
4244
},
4345
"cell_type": "code",
@@ -46,11 +48,9 @@
4648
"\n",
4749
"@numba.jit(nopython=True)\n",
4850
"def fibonacci_numba(b):\n",
49-
" n = 10\n",
5051
" b = b.copy()\n",
5152
" a = np.ones((), dtype=np.int32)\n",
52-
" # b = np.ones(1, dtype=np.int32)\n",
53-
" for _ in range(n):\n",
53+
" for _ in range(N_STEPS):\n",
5454
" a[()], b[()] = a[()] + b[()], a[()]\n",
5555
" return a"
5656
],
@@ -61,13 +61,12 @@
6161
{
6262
"metadata": {
6363
"ExecuteTime": {
64-
"end_time": "2025-10-07T09:52:11.125118Z",
65-
"start_time": "2025-10-07T09:52:07.865156Z"
64+
"end_time": "2025-10-07T10:20:26.947566Z",
65+
"start_time": "2025-10-07T10:20:23.615573Z"
6666
}
6767
},
6868
"cell_type": "code",
6969
"source": [
70-
"# n = np.int64(10)\n",
7170
"b = np.ones((), dtype=np.int32)\n",
7271
"assert fibonacci_pytensor(b) == fibonacci_numba(b)\n",
7372
"assert fibonacci_pytensor_numba(b) == fibonacci_numba(b)"
@@ -79,8 +78,8 @@
7978
{
8079
"metadata": {
8180
"ExecuteTime": {
82-
"end_time": "2025-10-07T09:52:14.968797Z",
83-
"start_time": "2025-10-07T09:52:11.202449Z"
81+
"end_time": "2025-10-07T10:20:28.819553Z",
82+
"start_time": "2025-10-07T10:20:27.015141Z"
8483
}
8584
},
8685
"cell_type": "code",
@@ -91,7 +90,7 @@
9190
"name": "stdout",
9291
"output_type": "stream",
9392
"text": [
94-
"46.6 μs ± 4.15 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
93+
"2.22 ms ± 37.8 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
9594
]
9695
}
9796
],
@@ -100,8 +99,8 @@
10099
{
101100
"metadata": {
102101
"ExecuteTime": {
103-
"end_time": "2025-10-07T09:52:19.495645Z",
104-
"start_time": "2025-10-07T09:52:15.020419Z"
102+
"end_time": "2025-10-07T10:20:42.302214Z",
103+
"start_time": "2025-10-07T10:20:28.871240Z"
105104
}
106105
},
107106
"cell_type": "code",
@@ -112,7 +111,7 @@
112111
"name": "stdout",
113112
"output_type": "stream",
114113
"text": [
115-
"5.53 μs ± 114 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
114+
"165 μs ± 468 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
116115
]
117116
}
118117
],
@@ -121,8 +120,8 @@
121120
{
122121
"metadata": {
123122
"ExecuteTime": {
124-
"end_time": "2025-10-07T09:52:22.257300Z",
125-
"start_time": "2025-10-07T09:52:19.547530Z"
123+
"end_time": "2025-10-07T10:20:55.256172Z",
124+
"start_time": "2025-10-07T10:20:42.355007Z"
126125
}
127126
},
128127
"cell_type": "code",
@@ -133,7 +132,7 @@
133132
"name": "stdout",
134133
"output_type": "stream",
135134
"text": [
136-
"3.37 μs ± 258 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
135+
"159 μs ± 1.41 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
137136
]
138137
}
139138
],
@@ -142,8 +141,8 @@
142141
{
143142
"metadata": {
144143
"ExecuteTime": {
145-
"end_time": "2025-10-07T09:52:32.320998Z",
146-
"start_time": "2025-10-07T09:52:22.315343Z"
144+
"end_time": "2025-10-07T10:20:57.954364Z",
145+
"start_time": "2025-10-07T10:20:55.346865Z"
147146
}
148147
},
149148
"cell_type": "code",
@@ -154,7 +153,7 @@
154153
"name": "stdout",
155154
"output_type": "stream",
156155
"text": [
157-
"1.19 μs ± 175 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
156+
"3.2 μs ± 19.2 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
158157
]
159158
}
160159
],
@@ -163,8 +162,8 @@
163162
{
164163
"metadata": {
165164
"ExecuteTime": {
166-
"end_time": "2025-10-07T09:52:32.381905Z",
167-
"start_time": "2025-10-07T09:52:32.372905Z"
165+
"end_time": "2025-10-07T10:20:58.015849Z",
166+
"start_time": "2025-10-07T10:20:58.006831Z"
168167
}
169168
},
170169
"cell_type": "code",
@@ -177,7 +176,7 @@
177176
"text": [
178177
"Subtensor{i} [id A] <Scalar(int32, shape=())> v={0: [0]} 6\n",
179178
" ├─ Scan{scan_fn, while_loop=False, inplace=all}.0 [id B] <Vector(int32, shape=(?,))> d={0: [1], 1: [2]} 5\n",
180-
" │ ├─ 10 [id C] <Scalar(int8, shape=())>\n",
179+
" │ ├─ 1000 [id C] <Scalar(int16, shape=())>\n",
181180
" │ ├─ SetSubtensor{:stop} [id D] <Vector(int32, shape=(1,))> d={0: [0]} 4\n",
182181
" │ │ ├─ AllocEmpty{dtype='int32'} [id E] <Vector(int32, shape=(1,))> 3\n",
183182
" │ │ │ └─ 1 [id F] <Scalar(int64, shape=())>\n",
@@ -203,7 +202,7 @@
203202
{
204203
"data": {
205204
"text/plain": [
206-
"<ipykernel.iostream.OutStream at 0x7f8462e7bbe0>"
205+
"<ipykernel.iostream.OutStream at 0x7fa4fbfacdf0>"
207206
]
208207
},
209208
"execution_count": 8,
@@ -216,8 +215,8 @@
216215
{
217216
"metadata": {
218217
"ExecuteTime": {
219-
"end_time": "2025-10-07T09:52:32.429959Z",
220-
"start_time": "2025-10-07T09:52:32.426275Z"
218+
"end_time": "2025-10-07T10:20:58.063585Z",
219+
"start_time": "2025-10-07T10:20:58.059985Z"
221220
}
222221
},
223222
"cell_type": "code",
@@ -239,7 +238,7 @@
239238
" tensor_variable_3 = allocempty_1(tensor_constant_1)\n",
240239
" # SetSubtensor{:stop}(AllocEmpty{dtype='int32'}.0, [1], 1)\n",
241240
" tensor_variable_4 = set_subtensor_1(tensor_variable_3, tensor_constant_2, scalar_constant)\n",
242-
" # Scan{scan_fn, while_loop=False, inplace=all}(10, SetSubtensor{:stop}.0, SetSubtensor{:stop}.0)\n",
241+
" # Scan{scan_fn, while_loop=False, inplace=all}(1000, SetSubtensor{:stop}.0, SetSubtensor{:stop}.0)\n",
243242
" tensor_variable_5, tensor_variable_6 = scan(tensor_constant_3, tensor_variable_4, tensor_variable_2)\n",
244243
" # Subtensor{i}(Scan{scan_fn, while_loop=False, inplace=all}.0, 0)\n",
245244
" tensor_variable_7 = subtensor(tensor_variable_5, scalar_constant_1)\n",
@@ -252,8 +251,8 @@
252251
{
253252
"metadata": {
254253
"ExecuteTime": {
255-
"end_time": "2025-10-07T09:52:32.479422Z",
256-
"start_time": "2025-10-07T09:52:32.475763Z"
254+
"end_time": "2025-10-07T10:20:58.113352Z",
255+
"start_time": "2025-10-07T10:20:58.109693Z"
257256
}
258257
},
259258
"cell_type": "code",
@@ -278,8 +277,8 @@
278277
{
279278
"metadata": {
280279
"ExecuteTime": {
281-
"end_time": "2025-10-07T09:52:32.529805Z",
282-
"start_time": "2025-10-07T09:52:32.526392Z"
280+
"end_time": "2025-10-07T10:20:58.162062Z",
281+
"start_time": "2025-10-07T10:20:58.158525Z"
283282
}
284283
},
285284
"cell_type": "code",
@@ -305,8 +304,8 @@
305304
{
306305
"metadata": {
307306
"ExecuteTime": {
308-
"end_time": "2025-10-07T09:52:32.579642Z",
309-
"start_time": "2025-10-07T09:52:32.576246Z"
307+
"end_time": "2025-10-07T10:20:58.211349Z",
308+
"start_time": "2025-10-07T10:20:58.207479Z"
310309
}
311310
},
312311
"cell_type": "code",
@@ -363,8 +362,8 @@
363362
{
364363
"metadata": {
365364
"ExecuteTime": {
366-
"end_time": "2025-10-07T09:52:32.638631Z",
367-
"start_time": "2025-10-07T09:52:32.625594Z"
365+
"end_time": "2025-10-07T10:20:58.266498Z",
366+
"start_time": "2025-10-07T10:20:58.254087Z"
368367
}
369368
},
370369
"cell_type": "code",
@@ -451,7 +450,7 @@
451450
" b_expanded = dimshuffle(b)\n",
452451
" b_buf_set = set_subtensor(b_buf, b_expanded, np.int64(1))\n",
453452
"\n",
454-
" a_buf_updated, b_buf_updated = scan_fib(np.array(10, dtype=\"int64\"), a_buf_set, b_buf_set)\n",
453+
" a_buf_updated, b_buf_updated = scan_fib(np.array(N_STEPS, np.int64), a_buf_set, b_buf_set)\n",
455454
"\n",
456455
" res = subtensor(a_buf_updated, np.uint8(0))\n",
457456
"\n",
@@ -464,8 +463,8 @@
464463
{
465464
"metadata": {
466465
"ExecuteTime": {
467-
"end_time": "2025-10-07T09:52:35.990806Z",
468-
"start_time": "2025-10-07T09:52:32.676348Z"
466+
"end_time": "2025-10-07T10:21:01.526942Z",
467+
"start_time": "2025-10-07T10:20:58.310079Z"
469468
}
470469
},
471470
"cell_type": "code",
@@ -480,8 +479,8 @@
480479
{
481480
"metadata": {
482481
"ExecuteTime": {
483-
"end_time": "2025-10-07T09:52:37.812669Z",
484-
"start_time": "2025-10-07T09:52:36.040144Z"
482+
"end_time": "2025-10-07T10:21:06.041171Z",
483+
"start_time": "2025-10-07T10:21:01.578757Z"
485484
}
486485
},
487486
"cell_type": "code",
@@ -492,7 +491,7 @@
492491
"name": "stdout",
493492
"output_type": "stream",
494493
"text": [
495-
"2.12 μs ± 51.6 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
494+
"55 μs ± 756 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
496495
]
497496
}
498497
],
@@ -501,8 +500,8 @@
501500
{
502501
"metadata": {
503502
"ExecuteTime": {
504-
"end_time": "2025-10-07T09:52:37.864066Z",
505-
"start_time": "2025-10-07T09:52:37.861775Z"
503+
"end_time": "2025-10-07T10:21:06.095536Z",
504+
"start_time": "2025-10-07T10:21:06.093418Z"
506505
}
507506
},
508507
"cell_type": "code",

0 commit comments

Comments
 (0)