JAX 易踩坑指南Common Gotchas純函數(shù)、原地更新、jit 與動態(tài)形狀的完整避坑手冊【免費下載鏈接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more項目地址: https://gitcode.com/GitHub_Trending/ja/jaxJAX 是一套對 Python NumPy 數(shù)值程序進行可組合變換微分、向量化、JIT 編譯到 GPU/TPU的框架但它的變換與編譯機制只對滿足特定約束的程序生效。本文以 docs/notebooks/Common_Gotchas_in_JAX.md 為骨架結(jié)合當(dāng)前倉庫源碼如 array_methods.py、slicing.py、config.py逐一拆解 JAX 中最常踩的坑純函數(shù)約束、原地更新、類方法 jit、越界索引、非數(shù)組輸入、隨機數(shù)、控制流、動態(tài)形狀、NaN/Inf 調(diào)試與 64 位精度并給出可復(fù)制、可運行的解決方案。讀完本文你將能識別并繞開這些陷阱寫出既正確又高性能的 JAX 代碼。 純函數(shù)Pure functionsJAX 的變換jit、grad、vmap等與編譯只設(shè)計用于函數(shù)式純函數(shù)所有輸入數(shù)據(jù)通過函數(shù)參數(shù)傳入所有結(jié)果通過函數(shù)返回值輸出。純函數(shù)對相同的輸入總是返回相同的結(jié)果。副作用Side-effects第一次運行與緩存命中的差異考慮帶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 # 副作用在第一次運行時出現(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值被捕獲還是實時讀取g 0. def impure_uses_globals(x): return x g # JAX 在第一次運行時捕獲全局變量的值 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實現(xiàn)其 docstring 給出了完整的等價對照表x.at寫法等價的原地表達式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)用順序是實現(xiàn)定義的、可能在部分硬件上不確定。越界切片尺寸限制在jit代碼以及l(fā)ax.while_loop/lax.fori_loop內(nèi)部切片的大小不能是參數(shù)值的函數(shù)只能依賴參數(shù)形狀切片起始索引不受此限制詳見下文控制流一節(jié)。 使用jax.jit裝飾類方法大多數(shù)jax.jit示例針對獨立函數(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 # ---- 如何正確做到這一點 def calc(self, y): if self.mul: return self.x * y return y調(diào)用c CustomClass(2, True); c.calc(3)會報錯因為函數(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)點簡單、顯式且無需教 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)后會作為字典鍵進入 JIT 的內(nèi)部編譯緩存因此 JAX 假定其哈希hash(obj)、相等性obj1 obj2與對象同一性obj1 is obj2行為一致。自定義對象的默認__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é)點精確指定哪些組件作為靜態(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 注冊機制的底層實現(xiàn)可參見 tree_util.py 中的register_pytree_node系列函數(shù)。 越界索引Out-of-bounds indexingNumPy 越界索引會拋錯np.arange(10)[11]但 JAX 無法或極難在加速器上拋出運行時錯誤因此必須為越界索引選擇某種非報錯行為類似于無效浮點運算產(chǎn)生NaN索引更新如index_add、scatter 類原語越界索引處的更新被跳過索引讀取如 NumPy 索引、gather 類原語索引被鉗制clamp到數(shù)組邊界因為必須返回某個東西。例如下面的索引操作會返回數(shù)組最后一個值jnp.arange(10)[11]這與底層GatherScatterMode的默認行為一致。查看 slicing.py 中GatherScatterMode的定義CLIP把索引鉗到最近的界內(nèi)值保證要 gather 的整個窗口都在界內(nèi)FILL_OR_DROPgather 時若窗口任何部分越界則整個窗口用常量填充scatter 時若窗口任何部分越界則整個窗口丟棄PROMISE_IN_BOUNDS用戶承諾索引在界內(nèi)不做額外檢查。當(dāng)前 XLA 實現(xiàn)下越界 gather 會被鉗制、越界 scatter 會被丟棄索引越界時梯度不正確。.at默認使用promise_in_bounds語義mode參數(shù)缺省值映射關(guān)系見 array_methods.py。用.at[...].get()精細控制越界行為如果需要對越界索引進行更精細的控制可用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默認get 鉗制 / set、add 等丟棄、clip鉗制、drop丟棄、filldrop的別名對get()可用fill_value指定返回值。另有wrap_negative_indices默認 True負索引表示從數(shù)組末尾數(shù)、indices_are_sorted與unique_indices提示實現(xiàn)索引已排序/唯一部分后端可優(yōu)化執(zhí)行若聲明與實際不符則輸出未定義。fill_value默認對非精確類型為NaN、有符號類型為最大負值、無符號類型為最大正值、布爾為True。兩個連帶影響由于索引讀取的鉗制行為jnp.nanargmin、jnp.nanargmax對全 NaN 切片返回 -1而 NumPy 會拋錯。上述兩種行為更新跳過 vs 讀取鉗制互不為逆運算因此反向模式自動微分把索引更新轉(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)作獨立的 JAX 變量被單獨處理并推送到設(shè)備。用make_jaxpr可以直觀看到make_jaxpr(permissive_sum)(x)每個 list 元素都被當(dāng)作獨立輸入處理追蹤與編譯開銷隨 list 長度線性增長。為避免此類意外JAX 不隱式轉(zhuǎn)換 list/tuple 為數(shù)組。若確實要向 JAX 函數(shù)傳 tuple/list請先顯式轉(zhuǎn)成數(shù)組jnp.sum(jnp.array(x)) 隨機數(shù)Random numbersJAX 的偽隨機數(shù)生成與 NumPy 有本質(zhì)區(qū)別。NumPy 使用隱式、全局、有狀態(tài)的隨機狀態(tài)JAX 采用顯式、無狀態(tài)的 PRNG key體系每個隨機操作都接收一個key參數(shù)并消費它新 key 通過random.split派生。快速上手可參考 docs/101/random.md 與 docs/101/index.rst 對應(yīng)的教程底層實現(xiàn)在 jax/_src/random/core.pykey、split、fold_in等核心函數(shù)默認使用 Threefry 計數(shù)器模式算法jax_default_prng_impl配置項默認threefry2x32見 config.py。key random.key(0) # 生成一個 PRNG key key, subkey random.split(key) # 分裂出新 key x random.uniform(subkey, (1000,))務(wù)必遵守每條隨機路徑使用獨立 key、用后即 split的規(guī)則避免可復(fù)現(xiàn)性被破壞。 控制流Control flow控制流細節(jié)已從本文移入專門的指南 docs/201/control-flow.mdjit對 Python 控制流與邏輯運算符的使用施加了約束需要用lax.cond、lax.while_loop、lax.fori_loop、lax.scan等結(jié)構(gòu)化控制流原語來表達依賴數(shù)據(jù)值的分支與循環(huán)。 動態(tài)形狀Dynamic shapes用于jax.jit、jax.vmap、jax.grad等變換的 JAX 代碼要求所有輸出數(shù)組與中間數(shù)組具有靜態(tài)形狀即形狀不能依賴其他數(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默認False。給每個操作添加 NaN 檢查當(dāng)在 jit 編譯計算輸出中檢測到 NaN 時回退到未編譯版本以更精確地定位產(chǎn)生 NaN 的操作jax_debug_infs默認False。與上同理針對 Inf 檢查。詳細使用方式見 docs/debugging/flags.md。 雙精度64-bit precisionJAX 默認強制單精度以緩解 NumPy API 將操作數(shù)激進提升到double的傾向。這對許多機器學(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 替你運行 absl 解析import jax if __name__ __main__: # 調(diào)用 jax.config.config_with_absl() 并執(zhí)行 absl 解析 jax.config.parse_flags_with_absl()注意方式 24 對 JAX 的任意配置選項都適用。確認 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 中定義的布爾配置項默認False且被標(biāo)記為include_in_jit_keyTrue、include_in_trace_contextTrue即它會參與 JIT 編譯鍵與追蹤上下文因此必須在啟動時盡早設(shè)置否則已緩存的編譯產(chǎn)物不會隨開關(guān)切換而更新。注意事項?? XLA 并非在所有后端都支持 64 位卷積 與 NumPy 的其他已知分歧jax.numpy盡力復(fù)刻 numpy API但存在一些 API 分歧的邊界情況除上文各節(jié)外已知分歧還包括類型提升規(guī)則二元運算中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)在浮點與整數(shù)類型之間轉(zhuǎn)換極端值時。次正規(guī)數(shù)subnormal刷新為零在部分后端上JAX 對次正規(guī)浮點數(shù)采用 flush-to-zero 語義 import jax.numpy as jnp subnormal jnp.float32(1E-45) subnormal # 次正規(guī)數(shù)本身可表示 Array(1.e-45, dtypefloat32) subnormal 0 # 但在運算內(nèi)被刷新為零 Array(0., dtypefloat32)次正規(guī)數(shù)的詳細運算語義通常隨后端而異。教程中覆蓋的其他坑docs/201/control-flow.md講解如何在jit對 Python 控制流與邏輯運算符的約束下工作docs/stateful-computations.md鑒于 JAX 變換只能作用于純函數(shù)該文給出在 JAX 程序中正確管理狀態(tài)的建議。Fin.如果本文沒有覆蓋到讓你抓狂的問題歡迎反饋以便擴充這份入門避坑指南。【免費下載鏈接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more項目地址: https://gitcode.com/GitHub_Trending/ja/jax創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考