Update translation: lectures/numpy_vs_numba_vs_jax.md

Author: mmcky

lectures/numpy_vs_numba_vs_jax.md

@@ -21,6 +21,8 @@ translation:
     Sequential operations: عملیات ترتیبی
     Sequential operations::Numba Version: نسخه Numba
     Sequential operations::JAX Version: نسخه JAX
+    Sequential operations::JAX Version::First Attempt: تلاش اول
+    Sequential operations::JAX Version::Second Attempt: تلاش دوم
     Sequential operations::Summary: خلاصه
     Overall recommendations: توصیه‌های کلی
 ---
@@ -137,33 +139,34 @@ m = -np.inf
 for x in grid:
     for y in grid:
         z = f(x, y)
-        if z > m:
-            m = z
+        m = max(m, z)
 ```
 
 ### برداری‌سازی NumPy
 
-اگر به برداری‌سازی به سبک NumPy تغییر دهیم، می‌توانیم از یک شبکه بسیار بزرگتر استفاده کنیم و کد نسبتاً سریع اجرا می‌شود.
+بیایید به NumPy تغییر دهیم و از یک شبکه بزرگتر استفاده کنیم
 
 در اینجا از `np.meshgrid` برای ایجاد شبکه‌های ورودی دوبعدی `x` و `y` استفاده می‌کنیم به گونه‌ای که `f(x, y)` تمام ارزیابی‌ها را روی شبکه حاصلضرب تولید می‌کند.
 
-(این استراتژی به Matlab بازمی‌گردد.)
-
 ```{code-cell} ipython3
+# Large grid
 grid = np.linspace(-3, 3, 3_000)
-x, y = np.meshgrid(grid, grid)
+
+x, y = np.meshgrid(grid, grid)    # MATLAB style meshgrid
 
 with qe.Timer():
     z_max_numpy = np.max(f(x, y))
-
-print(f"NumPy result: {z_max_numpy:.6f}")
 ```
 
 در نسخه برداری شده، تمام حلقه‌ها در کد کامپایل شده انجام می‌شوند.
 
-علاوه بر این، NumPy از چندنخی ضمنی استفاده می‌کند، به طوری که حداقل مقداری موازی‌سازی رخ می‌دهد.
+استفاده از `meshgrid` به ما امکان می‌دهد حلقه for تودرتو را تکرار کنیم.
 
-(موازی‌سازی نمی‌تواند بسیار کارآمد باشد زیرا فایل باینری قبل از اینکه اندازه آرایه‌های `x` و `y` را ببیند کامپایل می‌شود.)
+خروجی باید نزدیک به یک باشد:
+
+```{code-cell} ipython3
+print(f"NumPy result: {z_max_numpy:.6f}")
+```
 
 ### مقایسه با Numba
 
@@ -188,8 +191,6 @@ grid = np.linspace(-3, 3, 3_000)
 with qe.Timer():
     # First run
     z_max_numba = compute_max_numba(grid)
-
-print(f"Numba result: {z_max_numba:.6f}")
 ```
 
 بیایید دوباره اجرا کنیم تا زمان کامپایل حذف شود.
@@ -232,8 +233,6 @@ def compute_max_numba_parallel(grid):
 with qe.Timer():
     # First run
     z_max_parallel = compute_max_numba_parallel(grid)
-
-print(f"Numba result: {z_max_parallel:.6f}")
 ```
 
 در اینجا زمان‌بندی برای نسخه از پیش کامپایل شده آمده است.
@@ -244,27 +243,30 @@ with qe.Timer():
     compute_max_numba_parallel(grid)
 ```
 
-اگر چندین هسته دارید، باید حداقل برخی مزایا را از موازی‌سازی در اینجا ببینید.
+اگر چندین هسته دارید، باید مزایایی از موازی‌سازی در اینجا ببینید.
 
-برای دستگاه‌های قدرتمندتر و اندازه‌های شبکه بزرگتر، موازی‌سازی می‌تواند افزایش سرعت قابل توجهی ایجاد کند، حتی روی CPU.
+بیایید مطمئن شویم که نتیجه صحیح را به دست می‌آوریم (نزدیک به یک):
 
-### کد برداری شده با JAX
+```{code-cell} ipython3
+print(f"Numba result: {z_max_parallel:.6f}")
+```
+
+برای دستگاه‌های قدرتمند و اندازه‌های شبکه بزرگتر، موازی‌سازی می‌تواند افزایش سرعت مفیدی ایجاد کند، حتی روی CPU.
 
-در ظاهر، کد برداری شده در JAX شبیه به کد NumPy است.
+### کد برداری شده با JAX
 
-اما تفاوت‌هایی نیز وجود دارد که در اینجا آنها را برجسته می‌کنیم.
+بیایید رویکرد برداری شده NumPy را با JAX تکرار کنیم.
 
 بیایید با تابع شروع کنیم که `np` را به `jnp` تغییر می‌دهد و `jax.jit` را اضافه می‌کند.
 
-
 ```{code-cell} ipython3
 @jax.jit
 def f(x, y):
     return jnp.cos(x**2 + y**2) / (1 + x**2 + y**2)
 
 ```
 
-همانند NumPy، برای به دست آوردن شکل درست و محاسبه حلقه `for` تودرتوی صحیح، می‌توانیم از عملیات `meshgrid` طراحی شده برای این منظور استفاده کنیم:
+از رویکرد meshgrid به سبک NumPy استفاده می‌کنیم:
 
 ```{code-cell} ipython3
 grid = jnp.linspace(-3, 3, 3_000)
@@ -321,68 +323,37 @@ x_mesh.nbytes + y_mesh.nbytes
 
 در اینجا نحوه اعمال آن به مسئله ما آمده است.
 
-```{code-cell} ipython3
-# f را تنظیم کنید تا f(x, y) را در هر x برای هر y داده شده محاسبه کند
-f_vec_x = lambda y: f(grid, y)
-# یک تابع دوم ایجاد کنید که این عملیات را روی تمام y برداری کند
-f_vec = jax.vmap(f_vec_x)
-```
-
-اکنون `f_vec` هنگام فراخوانی با آرایه تخت `grid`، `f(x,y)` را در هر `x,y` محاسبه می‌کند.
-
-بیایید زمان‌بندی را ببینیم:
-
-```{code-cell} ipython3
-with qe.Timer():
-    z_max = jnp.max(f_vec(grid))
-    z_max.block_until_ready()
-
-print(f"JAX vmap v1 result: {z_max:.6f}")
-```
-
-```{code-cell} ipython3
-with qe.Timer():
-    z_max = jnp.max(f_vec(grid))
-    z_max.block_until_ready()
-```
-
-با اجتناب از آرایه‌های ورودی بزرگ `x_mesh` و `y_mesh`، این نسخه `vmap` از حافظه بسیار کمتری با زمان اجرای مشابه استفاده می‌کند.
-
-اما هنوز برخی بهره‌های سرعت را از دست می‌دهیم.
-
-کد فوق آرایه دوبعدی کامل `f(x,y)` را محاسبه می‌کند و سپس max را می‌گیرد.
-
-علاوه بر این، فراخوانی `jnp.max` خارج از تابع JIT-کامپایل شده `f` قرار دارد، بنابراین کامپایلر نمی‌تواند این عملیات را در یک kernel واحد ادغام کند.
-
-می‌توانیم هر دو مشکل را با انتقال max به داخل و پوشاندن همه چیز در یک `@jax.jit` واحد برطرف کنیم:
-
 ```{code-cell} ipython3
 @jax.jit
 def compute_max_vmap(grid):
-    # یک تابع بسازید که حداکثر را در امتداد هر سطر بگیرد
+    # Construct a function that takes the max over all x for given y
     f_vec_x_max = lambda y: jnp.max(f(grid, y))
-    # تابع را برداری کنید تا بتوانیم روی تمام سطرها همزمان فراخوانی کنیم
+    # Vectorize the function so we can call on all y simultaneously
     f_vec_max = jax.vmap(f_vec_x_max)
-    # تابع برداری شده را فراخوانی کنید و حداکثر را بگیرید
-    return jnp.max(f_vec_max(grid))
+    # Compute the max across x at every y
+    maxes = f_vec_max(grid)
+    # Compute the max of the maxes and return
+    return jnp.max(maxes)
 ```
 
-در اینجا
-
-* `f_vec_x_max` حداکثر را در امتداد هر سطر داده شده محاسبه می‌کند
-* `f_vec_max` یک نسخه برداری شده است که می‌تواند حداکثر تمام سطرها را به صورت موازی محاسبه کند.
+توجه کنید که هرگز
 
-ما این تابع را روی تمام سطرها اعمال می‌کنیم و سپس حداکثر max های سطر را می‌گیریم.
+* شبکه دوبعدی `x_mesh`
+* شبکه دوبعدی `y_mesh` یا
+* آرایه دوبعدی `f(x,y)`
 
-چون max را به داخل منتقل می‌کنیم، هرگز آرایه دوبعدی کامل `f(x,y)` را نمی‌سازیم و حافظه بیشتری صرفه‌جویی می‌شود.
+را نمی‌سازیم.
 
 و چون همه چیز زیر یک `@jax.jit` واحد قرار دارد، کامپایلر می‌تواند تمام عملیات را در یک kernel بهینه ادغام کند.
 
 بیایید آن را امتحان کنیم.
 
 ```{code-cell} ipython3
 with qe.Timer():
-    z_max = compute_max_vmap(grid).block_until_ready()
+    # First run
+    z_max = compute_max_vmap(grid)
+    # Hold interpreter
+    z_max.block_until_ready()
 
 print(f"JAX vmap result: {z_max:.6f}")
 ```
@@ -391,7 +362,10 @@ print(f"JAX vmap result: {z_max:.6f}")
 
 ```{code-cell} ipython3
 with qe.Timer():
-    z_max = compute_max_vmap(grid).block_until_ready()
+    # Second run
+    z_max = compute_max_vmap(grid)
+    # Hold interpreter
+    z_max.block_until_ready()
 ```
 
 ### خلاصه
@@ -448,13 +422,15 @@ with qe.Timer():
 
 Numba این عملیات ترتیبی را به طور بسیار کارآمد مدیریت می‌کند.
 
-توجه کنید که اجرای دوم پس از تکمیل کامپایل JIT به طور قابل توجهی سریعتر است.
+### نسخه JAX
+
+ما نمی‌توانیم مستقیماً `numba.jit` را با `jax.jit` جایگزین کنیم زیرا آرایه‌های JAX تغییرناپذیر هستند.
 
-کامپایل Numba معمولاً بسیار سریع است و عملکرد کد حاصل برای عملیات ترتیبی مانند این عالی است.
+اما می‌توانیم این عملیات را پیاده‌سازی کنیم.
 
-### نسخه JAX
+#### تلاش اول
 
-حالا بیایید یک نسخه JAX با استفاده از سینتکس `at[t].set` ایجاد کنیم که، همان‌طور که {ref}`در درس JAX بحث شد <jax_at_workaround>`، راه‌حلی برای آرایه‌های تغییرناپذیر فراهم می‌کند.
+در اینجا یک راه‌حل با استفاده از سینتکس `at[t].set` ارائه می‌شود که {ref}`در درس JAX بحث شد <jax_at_workaround>`.
 
 ما از `lax.fori_loop` استفاده می‌کنیم که نسخه‌ای از حلقه for است که می‌تواند توسط XLA کامپایل شود.
 
@@ -477,7 +453,7 @@ def qm_jax_fori(x0, n, α=4.0):
 * ما `n` را ایستا نگه می‌داریم زیرا بر اندازه آرایه تأثیر می‌گذارد و از این رو JAX می‌خواهد روی مقدار آن در کد کامپایل شده تخصصی شود.
 * ما به CPU از طریق `device=cpu` متصل می‌مانیم زیرا این بار کاری ترتیبی از بسیاری عملیات کوچک تشکیل شده است که فرصت کمی برای موازی‌سازی GPU باقی می‌گذارد.
 
-اگرچه `at[t].set` در هر مرحله ظاهراً یک آرایه جدید ایجاد می‌کند، در داخل یک تابع کامپایل‌شده با JIT، کامپایلر تشخیص می‌دهد که آرایه قدیمی دیگر مورد نیاز نیست و به‌روزرسانی را در جا انجام می‌دهد.
+مهم: اگرچه `at[t].set` در هر مرحله ظاهراً یک آرایه جدید ایجاد می‌کند، در داخل یک تابع کامپایل‌شده با JIT، کامپایلر تشخیص می‌دهد که آرایه قدیمی دیگر مورد نیاز نیست و به‌روزرسانی را در جا انجام می‌دهد!
 
 بیایید آن را با همان پارامترها زمان‌بندی کنیم:
 
@@ -499,7 +475,9 @@ with qe.Timer():
     x_jax.block_until_ready()
 ```
 
-JAX نیز برای این عملیات ترتیبی کاملاً کارآمد است.
+JAX نیز برای این عملیات ترتیبی کاملاً کارآمد است!
+
+#### تلاش دوم
 
 روش دیگری برای پیاده‌سازی حلقه وجود دارد که از `lax.scan` استفاده می‌کند.
 
@@ -538,11 +516,11 @@ with qe.Timer():
     x_jax.block_until_ready()
 ```
 
-هم JAX و هم Numba عملکرد قوی پس از کامپایل ارائه می‌دهند.
+شگفت‌انگیز است که JAX نیز پس از کامپایل عملکرد قوی ارائه می‌دهد.
 
 ### خلاصه
 
-در حالی که هم Numba و هم JAX عملکرد قوی برای عملیات ترتیبی ارائه می‌دهند، *تفاوت‌های قابل توجهی در خوانایی کد و سهولت استفاده وجود دارد*.
+در حالی که هم Numba و هم JAX عملکرد قوی برای عملیات ترتیبی ارائه می‌دهند، تفاوت‌هایی در خوانایی کد و سهولت استفاده وجود دارد.
 
 نسخه Numba ساده و طبیعی برای خواندن است: ما به سادگی یک آرایه اختصاص می‌دهیم و آن را عنصر به عنصر با استفاده از یک حلقه استاندارد Python پر می‌کنیم.