JAX 易踩坑指南Common Gotchas純函數(shù)、原地更新、jit 與動態(tài)形狀的完整避坑手冊【免費(fèi)下載鏈接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more項(xiàng)目地址: https://gitcode.com/GitHub_Trending/ja/jaxJAX 是一套對 Python NumPy 數(shù)值程序進(jìn)行可組合變換微分、向量化、JIT 編譯到 GPU/TPU的框架但它的變換與編譯機(jī)制只對滿足特定約束的程序生效。本文以 docs/notebooks/Common_Gotchas_in_JAX.md 為骨架結(jié)合當(dāng)前倉庫源碼如 array_methods.py、slicing.py、config.py逐一拆解 JAX 中最常踩的坑純函數(shù)約束、原地更新、類方法 jit、越界索引、非數(shù)組輸入、隨機(jī)數(shù)、控制流、動態(tài)形狀、NaN/Inf 調(diào)試與 64 位精度并給出可復(fù)制、可運(yùn)行的解決方案。讀完本文你將能識別并繞開這些陷阱寫出既正確又高性能的 JAX 代碼。 純函數(shù)Pure functionsJAX 的變換jit、grad、vmap等與編譯只設(shè)計用于函數(shù)式純函數(shù)所有輸入數(shù)據(jù)通過函數(shù)參數(shù)傳入所有結(jié)果通過函數(shù)返回值輸出。純函數(shù)對相同的輸入總是返回相同的結(jié)果。副作用Side-effects第一次運(yùn)行與緩存命中的差異考慮帶print副作用的函數(shù)import numpy as np from jax import jit from jax import lax from jax import random import jax import jax.numpy as jnp def impure_print_side_effect(x): print(Executing function) # 這是副作用 return x # 副作用在第一次運(yùn)行時出現(xiàn) print (First call: , jit(impure_print_side_effect)(4.)) # 相同類型/形狀參數(shù)的后續(xù)調(diào)用可能不再顯示副作用命中編譯緩存 print (Second call: , jit(impure_print_side_effect)(5.)) # 當(dāng)參數(shù)類型或形狀改變時JAX 會重新執(zhí)行 Python 函數(shù) print (Third call, different type: , jit(impure_print_side_effect)(jnp.array([5.])))由于 JAX 在參數(shù)類型與形狀不變時直接復(fù)用緩存的編譯結(jié)果print這類副作用只在首次或參數(shù)簽名變化時出現(xiàn)。需要注意這些行為并不被 JAX 系統(tǒng)保證正確用法是只對純函數(shù)使用 JAX 變換。全局變量Globals值被捕獲還是實(shí)時讀取g 0. def impure_uses_globals(x): return x g # JAX 在第一次運(yùn)行時捕獲全局變量的值 print (First call: , jit(impure_uses_globals)(4.)) g 10. # 更新全局變量 # 后續(xù)相同簽名的調(diào)用可能靜默使用緩存值 print (Second call: , jit(impure_uses_globals)(5.)) # 當(dāng)參數(shù)類型/形狀變化導(dǎo)致重新執(zhí)行 Python 時才會讀到最新的全局值 print (Third call, different type: , jit(impure_uses_globals)(jnp.array([4.])))修改全局變量小心內(nèi)部 Traced 值泄漏g 0. def impure_saves_global(x): global g g x return x # JAX 用特殊的 Traced 值執(zhí)行一次變換后的函數(shù) print (First call: , jit(impure_saves_global)(4.)) print (Saved global: , g) # 全局變量 g 被寫入了 JAX 內(nèi)部 Traced 值從源碼結(jié)構(gòu)看這正是 JAX 追蹤tracing模型的必然結(jié)果變換期間參數(shù)被替換為Traced抽象值寫入外部狀態(tài)的代碼會把內(nèi)部表示泄漏到 Python 對象中。正確的做法是把所有狀態(tài)當(dāng)作函數(shù)參數(shù)顯式傳入。內(nèi)部狀態(tài)是允許的只要不讀寫外部狀態(tài)一個 Python 函數(shù)只要不讀寫外部狀態(tài)即使內(nèi)部使用了有狀態(tài)對象也仍然是純函數(shù)def pure_uses_internal_state(x): state dict(even0, odd0) for i in range(10): state[even if i % 2 0 else odd] x return state[even] state[odd] print(jit(pure_uses_internal_state)(5.))不要在 jit 或控制流中使用迭代器迭代器是帶狀態(tài)的 Python 對象靠內(nèi)部狀態(tài)取下一個元素與 JAX 的函數(shù)式模型不兼容。在jit或任何控制流原語中使用迭代器大部分會直接報錯有些會靜默產(chǎn)生意外結(jié)果import jax.numpy as jnp from jax import make_jaxpr # lax.fori_loop直接用數(shù)組索引沒問題 array jnp.arange(10) print(lax.fori_loop(0, 10, lambda i,x: xarray[i], 0)) # 期望 45 iterator iter(range(10)) print(lax.fori_loop(0, 10, lambda i,x: xnext(iterator), 0)) # 意外結(jié)果 0 # lax.scan迭代器作為 elems 會拋錯 def func11(arr, extra): ones jnp.ones(arr.shape) def body(carry, aelems): ae1, ae2 aelems return (carry ae1 * ae2 extra, carry) return lax.scan(body, 0., (arr, ones)) make_jaxpr(func11)(jnp.arange(16), 5.) # make_jaxpr(func11)(iter(range(16)), 5.) # 拋錯 # lax.cond迭代器作為 operand 會拋錯 array_operand jnp.array([0.]) lax.cond(True, lambda x: x1, lambda x: x-1, array_operand) iter_operand iter(range(10)) # lax.cond(True, lambda x: next(x)1, lambda x: next(x)-1, iter_operand) # 拋錯lax.fori_loop、lax.scan、lax.cond的簽名與語義可參見倉庫源碼 lax/control_flow/loops.py 與 lax/control_flow/conditionals.py 相關(guān)定義。迭代器場景下應(yīng)改用數(shù)組或jnp.arange之類的無狀態(tài)數(shù)據(jù)結(jié)構(gòu)。 原地更新In-place updatesNumPy 中常見的原地索引更新numpy_array np.zeros((3,3), dtypenp.float32) print(original array:) print(numpy_array) # 原地、可變更新 numpy_array[1, :] 1.0 print(updated array:) print(numpy_array)而jax.Array禁止下標(biāo)賦值%xmode Minimal jax_array jnp.zeros((3,3), dtypejnp.float32) # 對 JAX 數(shù)組做原地更新會直接報錯 jax_array[1, :] 1.0__iadd__的差異重綁定而非原地修改jax_array jnp.array([10, 20]) jax_array_new jax_array jax_array_new 10 print(jax_array_new) # jax_array_new 被重綁定到新值 [20, 30]但... print(jax_array) # 原始數(shù)組保持 [10, 20] 不變 numpy_array np.array([10, 20]) numpy_array_new numpy_array numpy_array_new 10 print(numpy_array_new) # numpy_array_new is numpy_array被原地更新 print(numpy_array) # 兩者都是 [20, 30]原因在于NumPy 的__iadd__執(zhí)行原地修改而jax.Array不定義__iadd__Python 把jax_array_new 10當(dāng)作jax_array_new jax_array_new 10的語法糖只發(fā)生變量重綁定不修改任何數(shù)組。允許原地修改變量會讓程序分析與變換變得困難JAX 要求程序是純函數(shù)因此改用函數(shù)式數(shù)組更新。函數(shù)式數(shù)組更新x.at[idx].set(y)JAX 通過數(shù)組的.at屬性提供函數(shù)式純函數(shù)的索引更新。上面的更新可改寫為jax_array jnp.zeros((3,3), dtypejnp.float32) updated_array jax_array.at[1, :].set(1.0) print(updated array:\n, updated_array)與 NumPy 版本不同JAX 的數(shù)組更新函數(shù)**就地外out-of-place**操作返回新數(shù)組原始數(shù)組不被修改。print(original array unchanged:\n, jax_array)不過在jit 編譯代碼內(nèi)部如果x.at[idx].set(y)的輸入x之后不再被復(fù)用編譯器會自動把該數(shù)組更新優(yōu)化為原地執(zhí)行——這正是函數(shù)式寫法兼顧正確性與性能的關(guān)鍵。更多索引更新操作索引更新不只覆蓋數(shù)值還可以做索引加法等操作print(original array:) jax_array jnp.ones((5, 6)) print(jax_array) new_jax_array jax_array.at[::2, 3:].add(7.) print(new array post-addition:) print(new_jax_array)從源碼看.at由 array_methods.py 中的_IndexUpdateHelper實(shí)現(xiàn)其 docstring 給出了完整的等價對照表x.at寫法等價的原地表達(dá)式x x.at[idx].set(y)x[idx] yx x.at[idx].add(y)x[idx] yx x.at[idx].subtract(y)x[idx] - yx x.at[idx].multiply(y)x[idx] * yx x.at[idx].divide(y)x[idx] / yx x.at[idx].power(y)x[idx] ** yx x.at[idx].min(y)x[idx] minimum(x[idx], y)x x.at[idx].max(y)x[idx] maximum(x[idx], y)x x.at[idx].apply(ufunc)ufunc.at(x, idx)x x.at[idx].get()x x[idx]這些方法分別映射到底層 scatter 原語lax_slicing.scatter、scatter_add、scatter_sub、scatter_mul等見 array_methods.py。源碼 docstring 還提示與 NumPy 原地操作不同若多個索引指向同一位置所有更新都會被應(yīng)用NumPy 只保留最后一次且沖突更新的應(yīng)用順序是實(shí)現(xiàn)定義的、可能在部分硬件上不確定。越界切片尺寸限制在jit代碼以及l(fā)ax.while_loop/lax.fori_loop內(nèi)部切片的大小不能是參數(shù)值的函數(shù)只能依賴參數(shù)形狀切片起始索引不受此限制詳見下文控制流一節(jié)。 使用jax.jit裝飾類方法大多數(shù)jax.jit示例針對獨(dú)立函數(shù)裝飾類方法會引入額外問題?？匆粋€樸素寫法import jax.numpy as jnp from jax import jit class CustomClass: def __init__(self, x: jnp.ndarray, mul: bool): self.x x self.mul mul jit # ---- 如何正確做到這一點(diǎn) def calc(self, y): if self.mul: return self.x * y return y調(diào)用c CustomClass(2, True); c.calc(3)會報錯因?yàn)楹瘮?shù)第一個參數(shù)是self類型CustomClassJAX 不知道如何處理這種類型。文檔給出三種基本策略。策略 1JIT 編譯的輔助函數(shù)helper function最直接的方式是在類外定義一個可正常 JIT 的輔助函數(shù)class CustomClass: def __init__(self, x: jnp.ndarray, mul: bool): self.x x self.mul mul def calc(self, y): return _calc(self.mul, self.x, y) jit(static_argnums0) def _calc(mul, x, y): if mul: return x * y return yc CustomClass(2, True) print(c.calc(3))優(yōu)點(diǎn)簡單、顯式且無需教 JAX 如何處理CustomClass類型代價是方法邏輯被拆到了類外。策略 2把self標(biāo)記為靜態(tài)static用static_argnums標(biāo)記self為靜態(tài)參數(shù)但要小心意外結(jié)果class CustomClass: def __init__(self, x: jnp.ndarray, mul: bool): self.x x self.mul mul # 警告下面的例子是壞的別直接復(fù)制粘貼 jit(static_argnums0) def calc(self, y): if self.mul: return self.x * y return y首次調(diào)用c CustomClass(2, True); print(c.calc(3))不再報錯但陷阱在于首次調(diào)用后修改對象屬性后續(xù)調(diào)用可能返回錯誤結(jié)果c.mul False print(c.calc(3)) # 本應(yīng)打印 3原因?qū)ο蟊粯?biāo)記為靜態(tài)后會作為字典鍵進(jìn)入 JIT 的內(nèi)部編譯緩存因此 JAX 假定其哈希hash(obj)、相等性obj1 obj2與對象同一性obj1 is obj2行為一致。自定義對象的默認(rèn)__hash__是其對象 ID所以 JAX 無從得知對象被修改應(yīng)觸發(fā)重新編譯。部分解決方法是定義合適的__hash__與__eq__class CustomClass: def __init__(self, x: jnp.ndarray, mul: bool): self.x x self.mul mul jit(static_argnums0) def calc(self, y): if self.mul: return self.x * y return y def __hash__(self): return hash((self.x, self.mul)) def __eq__(self, other): return (isinstance(other, CustomClass) and (self.x, self.mul) (other.x, other.mul))只要永不修改對象這種方式能配合 JIT 與其他變換正常工作。對象一旦被修改作為哈希鍵使用會引發(fā)多種微妙問題——這正是可變?nèi)萜鱠ict、list不定義__hash__、而不可變?nèi)萜鱰uple定義的原因。若你的類依賴原地修改如方法內(nèi)self.attr ...對象就不是真正靜態(tài)的標(biāo)記為靜態(tài)會出問題——這時應(yīng)該用策略 3。策略 3把CustomClass注冊為 PyTree最靈活的做法是把類型注冊為自定義 PyTree 節(jié)點(diǎn)精確指定哪些組件作為靜態(tài)aux data、哪些作為動態(tài)childrenclass CustomClass: def __init__(self, x: jnp.ndarray, mul: bool): self.x x self.mul mul jit def calc(self, y): if self.mul: return self.x * y return y def _tree_flatten(self): children (self.x,) # 數(shù)組 / 動態(tài)值 aux_data {mul: self.mul} # 靜態(tài)值 return (children, aux_data) classmethod def _tree_unflatten(cls, aux_data, children): return cls(*children, **aux_data) from jax import tree_util tree_util.register_pytree_node(CustomClass, CustomClass._tree_flatten, CustomClass._tree_unflatten)該方案解決了前述所有問題c CustomClass(2, True) print(c.calc(3)) c.mul False # 修改會被檢測到 print(c.calc(3)) c CustomClass(jnp.array(2), True) # 不可哈希的 x 也能支持 print(c.calc(3))只要tree_flatten/tree_unflatten正確處理類的所有相關(guān)屬性就能不加任何特殊注解直接將該類型的對象作為 JIT 函數(shù)的參數(shù)。PyTree 注冊機(jī)制的底層實(shí)現(xiàn)可參見 tree_util.py 中的register_pytree_node系列函數(shù)。 越界索引Out-of-bounds indexingNumPy 越界索引會拋錯np.arange(10)[11]但 JAX 無法或極難在加速器上拋出運(yùn)行時錯誤因此必須為越界索引選擇某種非報錯行為類似于無效浮點(diǎn)運(yùn)算產(chǎn)生NaN索引更新如index_add、scatter 類原語越界索引處的更新被跳過索引讀取如 NumPy 索引、gather 類原語索引被鉗制clamp到數(shù)組邊界因?yàn)楸仨毞祷啬硞€東西。例如下面的索引操作會返回數(shù)組最后一個值jnp.arange(10)[11]這與底層GatherScatterMode的默認(rèn)行為一致。查看 slicing.py 中GatherScatterMode的定義CLIP把索引鉗到最近的界內(nèi)值保證要 gather 的整個窗口都在界內(nèi)FILL_OR_DROPgather 時若窗口任何部分越界則整個窗口用常量填充scatter 時若窗口任何部分越界則整個窗口丟棄PROMISE_IN_BOUNDS用戶承諾索引在界內(nèi)不做額外檢查。當(dāng)前 XLA 實(shí)現(xiàn)下越界 gather 會被鉗制、越界 scatter 會被丟棄索引越界時梯度不正確。.at默認(rèn)使用promise_in_bounds語義mode參數(shù)缺省值映射關(guān)系見 array_methods.py。用.at[...].get()精細(xì)控制越界行為如果需要對越界索引進(jìn)行更精細(xì)的控制可用ndarray.at的可選參數(shù)jnp.arange(10.0).at[11].get()jnp.arange(10.0).at[11].get(modefill, fill_valuejnp.nan)mode取值promise_in_bounds默認(rèn)get 鉗制 / set、add 等丟棄、clip鉗制、drop丟棄、filldrop的別名對get()可用fill_value指定返回值。另有wrap_negative_indices默認(rèn) True負(fù)索引表示從數(shù)組末尾數(shù)、indices_are_sorted與unique_indices提示實(shí)現(xiàn)索引已排序/唯一部分后端可優(yōu)化執(zhí)行若聲明與實(shí)際不符則輸出未定義。fill_value默認(rèn)對非精確類型為NaN、有符號類型為最大負(fù)值、無符號類型為最大正值、布爾為True。兩個連帶影響由于索引讀取的鉗制行為jnp.nanargmin、jnp.nanargmax對全 NaN 切片返回 -1而 NumPy 會拋錯。上述兩種行為更新跳過 vs 讀取鉗制互不為逆運(yùn)算因此反向模式自動微分把索引更新轉(zhuǎn)成索引讀取、反之亦然不會保持越界索引語義。最好把 JAX 中的越界索引視為未定義行為undefined behavior。 非數(shù)組輸入NumPy vs. JAXNumPy 通常樂意接受 Python list / tuple 作為 API 輸入np.sum([1, 2, 3])JAX 則相反通常會給出有用的報錯jnp.sum([1, 2, 3])這是刻意的設(shè)計決策把 list/tuple 傳給被追蹤traced的函數(shù)可能導(dǎo)致難以察覺的靜默性能退化。例如下面這個寬容版jnp.sumdef permissive_sum(x): return jnp.sum(jnp.array(x)) x list(range(10)) permissive_sum(x)輸出符合預(yù)期但掩蓋了性能問題在 JAX 的追蹤與 JIT 編譯模型里Python list/tuple 的每個元素都被當(dāng)作獨(dú)立的 JAX 變量被單獨(dú)處理并推送到設(shè)備。用make_jaxpr可以直觀看到make_jaxpr(permissive_sum)(x)每個 list 元素都被當(dāng)作獨(dú)立輸入處理追蹤與編譯開銷隨 list 長度線性增長。為避免此類意外JAX 不隱式轉(zhuǎn)換 list/tuple 為數(shù)組。若確實(shí)要向 JAX 函數(shù)傳 tuple/list請先顯式轉(zhuǎn)成數(shù)組jnp.sum(jnp.array(x)) 隨機(jī)數(shù)Random numbersJAX 的偽隨機(jī)數(shù)生成與 NumPy 有本質(zhì)區(qū)別。NumPy 使用隱式、全局、有狀態(tài)的隨機(jī)狀態(tài)JAX 采用顯式、無狀態(tài)的 PRNG key體系每個隨機(jī)操作都接收一個key參數(shù)并消費(fèi)它新 key 通過random.split派生?？焖偕鲜挚蓞⒖?docs/101/random.md 與 docs/101/index.rst 對應(yīng)的教程底層實(shí)現(xiàn)在 jax/_src/random/core.pykey、split、fold_in等核心函數(shù)默認(rèn)使用 Threefry 計數(shù)器模式算法jax_default_prng_impl配置項(xiàng)默認(rèn)threefry2x32見 config.py。key random.key(0) # 生成一個 PRNG key key, subkey random.split(key) # 分裂出新 key x random.uniform(subkey, (1000,))務(wù)必遵守每條隨機(jī)路徑使用獨(dú)立 key、用后即 split的規(guī)則避免可復(fù)現(xiàn)性被破壞。 控制流Control flow控制流細(xì)節(jié)已從本文移入專門的指南 docs/201/control-flow.mdjit對 Python 控制流與邏輯運(yùn)算符的使用施加了約束需要用lax.cond、lax.while_loop、lax.fori_loop、lax.scan等結(jié)構(gòu)化控制流原語來表達(dá)依賴數(shù)據(jù)值的分支與循環(huán)。 動態(tài)形狀Dynamic shapes用于jax.jit、jax.vmap、jax.grad等變換的 JAX 代碼要求所有輸出數(shù)組與中間數(shù)組具有靜態(tài)形狀即形狀不能依賴其他數(shù)組中的值。例如自己實(shí)現(xiàn)jnp.nansum時可能這樣寫def nansum(x): mask ~jnp.isnan(x) # 選擇非 NaN 值的布爾掩碼 x_without_nans x[mask] return x_without_nans.sum()在 JIT 之外它可以正常工作x jnp.array([1, 2, jnp.nan, 3, 4]) print(nansum(x))但對其應(yīng)用jax.jit或其他變換就會報錯jax.jit(nansum)(x)問題在于x_without_nans的尺寸依賴x中的值即它是動態(tài)的。JAX 中通常可以用其他手段繞開動態(tài)形狀數(shù)組例如用三參數(shù)形式j(luò)np.where把 NaN 替換為 0得到相同結(jié)果同時避免動態(tài)形狀jax.jit def nansum_2(x): mask ~jnp.isnan(x) # 選擇非 NaN 值的布爾掩碼 return jnp.where(mask, x, 0).sum() print(nansum_2(x))其他出現(xiàn)動態(tài)形狀數(shù)組的場景也可采用類似技巧。 調(diào)試 NaN 與 Inf使用jax_debug_nans與jax_debug_infs兩個 flag 定位函數(shù)與梯度中 NaN/Inf 的來源。它們定義于 config.pyjax_debug_nans默認(rèn)False。給每個操作添加 NaN 檢查當(dāng)在 jit 編譯計算輸出中檢測到 NaN 時回退到未編譯版本以更精確地定位產(chǎn)生 NaN 的操作jax_debug_infs默認(rèn)False。與上同理針對 Inf 檢查。詳細(xì)使用方式見 docs/debugging/flags.md。 雙精度64-bit precisionJAX 默認(rèn)強(qiáng)制單精度以緩解 NumPy API 將操作數(shù)激進(jìn)提升到double的傾向。這對許多機(jī)器學(xué)習(xí)應(yīng)用是期望行為但可能出乎你的意料x random.uniform(random.key(0), (1000,), dtypejnp.float64) x.dtype輸出仍是float32。要使用雙精度必須在**啟動時startup**設(shè)置jax_enable_x64配置變量有以下幾種方式設(shè)置環(huán)境變量JAX_ENABLE_X64True啟動時手動設(shè)置配置 flag# 注意這只能在啟動時生效 import jax jax.config.update(jax_enable_x64, True)用absl.app.run(main)解析命令行 flagsimport jax jax.config.config_with_absl()讓 JAX 替你運(yùn)行 absl 解析import jax if __name__ __main__: # 調(diào)用 jax.config.config_with_absl() 并執(zhí)行 absl 解析 jax.config.parse_flags_with_absl()注意方式 24 對 JAX 的任意配置選項(xiàng)都適用。確認(rèn) x64 已啟用import jax import jax.numpy as jnp from jax import random jax.config.update(jax_enable_x64, True) x random.uniform(random.key(0), (1000,), dtypejnp.float64) x.dtype # -- dtype(float64)從源碼看jax_enable_x64是 config.py 中定義的布爾配置項(xiàng)默認(rèn)False且被標(biāo)記為include_in_jit_keyTrue、include_in_trace_contextTrue即它會參與 JIT 編譯鍵與追蹤上下文因此必須在啟動時盡早設(shè)置否則已緩存的編譯產(chǎn)物不會隨開關(guān)切換而更新。注意事項(xiàng)?? XLA 并非在所有后端都支持 64 位卷積 與 NumPy 的其他已知分歧jax.numpy盡力復(fù)刻 numpy API但存在一些 API 分歧的邊界情況除上文各節(jié)外已知分歧還包括類型提升規(guī)則二元運(yùn)算中JAX 的類型提升規(guī)則與 NumPy 略有不同詳見 docs/101/type_promotion.rst。不安全類型轉(zhuǎn)換unsafe cast當(dāng)目標(biāo) dtype 無法表示輸入值時JAX 的行為可能依賴后端總體上可能與 NumPy 不同。NumPy 通過astype的casting參數(shù)控制結(jié)果JAX 不提供此類配置直接繼承 XLAConvertElementType的語義。例如 np.arange(254.0, 258.0).astype(uint8) array([254, 255, 0, 1], dtypeuint8) jnp.arange(254.0, 258.0).astype(uint8) Array([254, 255, 255, 255], dtypeuint8)這類不一致典型出現(xiàn)在浮點(diǎn)與整數(shù)類型之間轉(zhuǎn)換極端值時。次正規(guī)數(shù)subnormal刷新為零在部分后端上JAX 對次正規(guī)浮點(diǎn)數(shù)采用 flush-to-zero 語義 import jax.numpy as jnp subnormal jnp.float32(1E-45) subnormal # 次正規(guī)數(shù)本身可表示 Array(1.e-45, dtypefloat32) subnormal 0 # 但在運(yùn)算內(nèi)被刷新為零 Array(0., dtypefloat32)次正規(guī)數(shù)的詳細(xì)運(yùn)算語義通常隨后端而異。教程中覆蓋的其他坑docs/201/control-flow.md講解如何在jit對 Python 控制流與邏輯運(yùn)算符的約束下工作docs/stateful-computations.md鑒于 JAX 變換只能作用于純函數(shù)該文給出在 JAX 程序中正確管理狀態(tài)的建議。Fin.如果本文沒有覆蓋到讓你抓狂的問題歡迎反饋以便擴(kuò)充這份入門避坑指南?！久赓M(fèi)下載鏈接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more項(xiàng)目地址: https://gitcode.com/GitHub_Trending/ja/jax創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考