Update translation: lectures/numpy_vs_numba_vs_jax.md

Author: mmcky

lectures/numpy_vs_numba_vs_jax.md

@@ -36,23 +36,23 @@ translation:
 
 # NumPy در مقابل Numba در مقابل JAX
 
-در سخنرانی‌های قبلی، سه کتابخانه اصلی برای محاسبات علمی و عددی را بحث کردیم:
+در درس‌های قبلی، سه کتابخانه اصلی برای محاسبات علمی و عددی را بحث کردیم:
 
 * [NumPy](numpy)
 * [Numba](numba)
 * [JAX](jax_intro)
 
 کدام یک را باید در هر موقعیت استفاده کنیم؟
 
-این سخنرانی به آن سؤال پاسخ می‌دهد، حداقل تا حدی، با بحث در مورد برخی موارد استفاده.
+این درس به آن سؤال پاسخ می‌دهد، حداقل تا حدی، با بحث در مورد برخی موارد استفاده.
 
 قبل از شروع، توجه می‌کنیم که دو مورد اول یک جفت طبیعی هستند: NumPy و Numba به خوبی با هم کار می‌کنند.
 
 JAX، از سوی دیگر، به تنهایی می‌ایستد.
 
 هنگام بررسی هر رویکرد، نه تنها کارایی و رد پای حافظه، بلکه وضوح و سهولت استفاده را نیز در نظر خواهیم گرفت.
 
-علاوه بر آنچه در Anaconda موجود است، این سخنرانی به کتابخانه‌های زیر نیاز دارد:
+علاوه بر آنچه در Anaconda موجود است، این درس به کتابخانه‌های زیر نیاز دارد:
 
 ```{code-cell} ipython3
 ---
@@ -67,7 +67,6 @@ tags: [hide-output]
 ما از import های زیر استفاده خواهیم کرد.
 
 ```{code-cell} ipython3
-import random
 from functools import partial
 
 import numpy as np
@@ -455,15 +454,60 @@ Numba این عملیات ترتیبی را به طور بسیار کارآمد
 
 ### نسخه JAX
 
-حالا بیایید یک نسخه JAX با استفاده از `lax.scan` ایجاد کنیم:
+حالا بیایید یک نسخه JAX با استفاده از سینتکس `at[t].set` ایجاد کنیم که، همان‌طور که {ref}`در درس JAX بحث شد <jax_at_workaround>`، راه‌حلی برای آرایه‌های تغییرناپذیر فراهم می‌کند.
 
-(ما `n` را ایستا نگه می‌داریم زیرا بر اندازه آرایه تأثیر می‌گذارد و از این رو JAX می‌خواهد روی مقدار آن در کد کامپایل شده تخصصی شود.)
+ما از `lax.fori_loop` استفاده می‌کنیم که نسخه‌ای از حلقه for است که می‌تواند توسط XLA کامپایل شود.
 
 ```{code-cell} ipython3
 cpu = jax.devices("cpu")[0]
 
-@partial(jax.jit, static_argnames=('n',), device=cpu)
-def qm_jax(x0, n, α=4.0):
+@partial(jax.jit, static_argnames=("n",), device=cpu)
+def qm_jax_fori(x0, n, α=4.0):
+
+    x = jnp.empty(n + 1).at[0].set(x0)
+
+    def update(t, x):
+        return x.at[t + 1].set(α * x[t] * (1 - x[t]))
+
+    x = lax.fori_loop(0, n, update, x)
+    return x
+
+```
+
+* ما `n` را ایستا نگه می‌داریم زیرا بر اندازه آرایه تأثیر می‌گذارد و از این رو JAX می‌خواهد روی مقدار آن در کد کامپایل شده تخصصی شود.
+* ما به CPU از طریق `device=cpu` متصل می‌مانیم زیرا این بار کاری ترتیبی از بسیاری عملیات کوچک تشکیل شده است که فرصت کمی برای موازی‌سازی GPU باقی می‌گذارد.
+
+اگرچه `at[t].set` در هر مرحله ظاهراً یک آرایه جدید ایجاد می‌کند، در داخل یک تابع کامپایل‌شده با JIT، کامپایلر تشخیص می‌دهد که آرایه قدیمی دیگر مورد نیاز نیست و به‌روزرسانی را در جا انجام می‌دهد.
+
+بیایید آن را با همان پارامترها زمان‌بندی کنیم:
+
+```{code-cell} ipython3
+with qe.Timer():
+    # First run
+    x_jax = qm_jax_fori(0.1, n)
+    # Hold interpreter
+    x_jax.block_until_ready()
+```
+
+بیایید دوباره اجرا کنیم تا سربار کامپایل حذف شود:
+
+```{code-cell} ipython3
+with qe.Timer():
+    # Second run
+    x_jax = qm_jax_fori(0.1, n)
+    # Hold interpreter
+    x_jax.block_until_ready()
+```
+
+JAX نیز برای این عملیات ترتیبی کاملاً کارآمد است.
+
+روش دیگری برای پیاده‌سازی حلقه وجود دارد که از `lax.scan` استفاده می‌کند.
+
+این روش جایگزین، به طور قابل بحث، بیشتر با رویکرد تابعی JAX همسو است --- اگرچه سینتکس آن به خاطر سپردن دشواری دارد.
+
+```{code-cell} ipython3
+@partial(jax.jit, static_argnames=("n",), device=cpu)
+def qm_jax_scan(x0, n, α=4.0):
     def update(x, t):
         x_new = α * x * (1 - x)
         return x_new, x_new
@@ -474,16 +518,12 @@ def qm_jax(x0, n, α=4.0):
 
 این کد خواندن آسانی ندارد اما، در اصل، `lax.scan` به طور مکرر `update` را فراخوانی می‌کند و بازگشت‌های `x_new` را در یک آرایه جمع می‌کند.
 
-```{note}
-ما `device=cpu` را در decorator `jax.jit` مشخص می‌کنیم زیرا این محاسبه از بسیاری عملیات ترتیبی کوچک تشکیل شده است که فرصت کمی برای بهره‌برداری GPU از موازی‌سازی باقی می‌گذارد. در نتیجه، سربار راه‌اندازی kernel تمایل دارد روی GPU غالب شود و CPU را متناسب‌تر برای این بار کاری می‌کند.
-```
-
 بیایید آن را با همان پارامترها زمان‌بندی کنیم:
 
 ```{code-cell} ipython3
 with qe.Timer():
     # First run
-    x_jax = qm_jax(0.1, n)
+    x_jax = qm_jax_scan(0.1, n)
     # Hold interpreter
     x_jax.block_until_ready()
 ```
@@ -493,13 +533,11 @@ with qe.Timer():
 ```{code-cell} ipython3
 with qe.Timer():
     # Second run
-    x_jax = qm_jax(0.1, n)
+    x_jax = qm_jax_scan(0.1, n)
     # Hold interpreter
     x_jax.block_until_ready()
 ```
 
-JAX نیز برای این عملیات ترتیبی کاملاً کارآمد است.
-
 هم JAX و هم Numba عملکرد قوی پس از کامپایل ارائه می‌دهند.
 
 ### خلاصه
@@ -510,9 +548,9 @@ JAX نیز برای این عملیات ترتیبی کاملاً کارآمد 
 
 این دقیقاً نحوه تفکر اکثر برنامه‌نویسان در مورد الگوریتم است.
 
-نسخه JAX، از سوی دیگر، نیاز به استفاده از `lax.scan` دارد که به طور قابل توجهی کمتر شهودی است.
+نسخه‌های JAX، از سوی دیگر، نیاز به استفاده از `lax.fori_loop` یا `lax.scan` دارند که هر دو کمتر شهودی از یک حلقه استاندارد Python هستند.
 
-علاوه بر این، آرایه‌های تغییرناپذیر JAX به این معنی است که نمی‌توانیم به سادگی عناصر آرایه را در جا به‌روزرسانی کنیم و تکرار مستقیم الگوریتم مورد استفاده توسط Numba را سخت می‌کند.
+در حالی که سینتکس `at[t].set` در JAX به‌روزرسانی عنصر به عنصر را ممکن می‌سازد، کد کلی همچنان سخت‌تر از معادل Numba برای خواندن است.
 
 برای این نوع عملیات ترتیبی، Numba برنده واضح از نظر وضوح کد و سهولت پیاده‌سازی است.
 
@@ -532,11 +570,12 @@ JAX نیز برای این عملیات ترتیبی کاملاً کارآمد 
 
 کد طبیعی و خوانا است --- صرفاً یک حلقه پایتون با یک decorator --- و کارایی آن عالی است.
 
-JAX می‌تواند مسائل ترتیبی را از طریق `lax.scan` مدیریت کند، اما نحو آن کمتر شهودی است و برای کارهای کاملاً ترتیبی، بهره‌وری اضافی ناچیز است.
-
-با این حال، `lax.scan` یک مزیت مهم دارد: از مشتق‌گیری خودکار در طول حلقه پشتیبانی می‌کند، که Numba قادر به انجام آن نیست.
+JAX می‌تواند مسائل ترتیبی را از طریق `lax.fori_loop` یا `lax.scan` مدیریت کند، اما نحو آن کمتر شهودی است.
 
+```{note}
+یک مزیت مهم `lax.fori_loop` و `lax.scan` این است که از مشتق‌گیری خودکار در طول حلقه پشتیبانی می‌کنند، که Numba قادر به انجام آن نیست.
 اگر نیاز دارید از طریق یک محاسبه ترتیبی مشتق بگیرید (مثلاً محاسبه حساسیت‌های یک مسیر نسبت به پارامترهای مدل)، JAX علی‌رغم نحو کمتر طبیعی‌اش، انتخاب بهتری است.
+```
 
 در عمل، بسیاری از مسائل ترکیبی از هر دو الگو هستند.