JAX 是 TensorFlow 和 PyTorch 的新競爭對手。 JAX 強調(diào)簡單性而不犧牲速度和可擴展性。由于 JAX 需要更少的樣板代碼，因此程序更短、更接近數(shù)學(xué)，因此更容易理解。
長話短說：
-
使用 import jax.numpy 訪問 NumPy 函數(shù)，使用 import jax.scipy 訪問 SciPy 函數(shù)。 -
通過使用 @jax.jit 進行裝飾，可以加快即時編譯速度。 -
使用 jax.grad 求導(dǎo)。 -
使用 jax.vmap 進行矢量化，并使用 jax.pmap 進行跨設(shè)備并行化。
函數(shù)式編程
JAX 遵循函數(shù)式編程哲學(xué)。這意味著您的函數(shù)必須是獨立的或純粹的：不允許有副作用。本質(zhì)上，純函數(shù)看起來像數(shù)學(xué)函數(shù)（圖 1）。有輸入進來，有東西出來，但與外界沒有溝通。
例子#1
以下代碼片段是一個非功能純的示例。
import?jax.numpy?as?jnp
bias?=?jnp.array(0)
def?impure_example(x):
???total?=?x?+?bias
???return?total
注意 impure_example 之外的偏差。在編譯期間（見下文），偏差可能會被緩存，因此不再反映偏差的變化。
例子#2
這是一個pure的例子。
def?pure_example(x,?weights,?bias):
???activation?=?weights?@?x?+?bias
???return?activation
在這里，pure_example 是獨立的：所有參數(shù)都作為參數(shù)傳遞。
確定性采樣器
在計算機中，不存在真正的隨機性。相反，NumPy 和 TensorFlow 等庫會跟蹤偽隨機數(shù)狀態(tài)來生成“隨機”樣本。
函數(shù)式編程的直接后果是隨機函數(shù)的工作方式不同。由于不再允許全局狀態(tài)，因此每次采樣隨機數(shù)時都需要顯式傳入偽隨機數(shù)生成器 (PRNG) 密鑰
import?jax
key?=?jax.random.PRNGKey(42)
u?=?jax.random.uniform(key)
此外，您有責任為任何后續(xù)調(diào)用推進“隨機狀態(tài)”。
key?=?jax.random.PRNGKey(43)
#?Split?off?and?consume?subkey.
key,?subkey?=?jax.random.split(key)
u?=?jax.random.uniform(subkey)
#?Split?off?and?consume?second?subkey.
key,?subkey?=?jax.random.split(key)
u?=?jax.random.uniform(subkey)
..
jit
您可以通過即時編譯 JAX 指令來加快代碼速度。例如，要編譯縮放指數(shù)線性單位 (SELU) 函數(shù)，請使用 jax.numpy 中的 NumPy 函數(shù)并將 jax.jit 裝飾器添加到該函數(shù)，如下所示：
from?jax?import?jit
@jit
def?selu(x,?α=1.67,?λ=1.05):
?return?λ?*?jnp.where(x?>?0,?x,?α?*?jnp.exp(x)?-?α)
JAX 會跟蹤您的指令并將其轉(zhuǎn)換為 jaxpr。這使得加速線性代數(shù) (XLA) 編譯器能夠為您的加速器生成非常高效的優(yōu)化代碼。
gard
JAX 最強大的功能之一是您可以輕松獲取 gard。使用 jax.grad，您可以定義一個新函數(shù)，即符號導(dǎo)數(shù)。
from?jax?import?grad
def?f(x):
???return?x?+?0.5?*?x**2
df_dx?=?grad(f)
d2f_dx2?=?grad(grad(f))
正如您在示例中看到的，您不僅限于一階導(dǎo)數(shù)。您可以通過簡單地按順序鏈接 grad 函數(shù) n 次來獲取 n 階導(dǎo)數(shù)。
vmap 和 pmap
矩陣乘法使所有批次尺寸正確需要非常細心。 JAX 的矢量化映射函數(shù) vmap 通過對函數(shù)進行矢量化來減輕這種負擔。基本上，每個按元素應(yīng)用函數(shù) f 的代碼塊都是由 vmap 替換的候選者。讓我們看一個例子。
計算線性函數(shù)：
def?linear(x):
?return?weights?@?x
在一批示例 [x?, x2,..] 中，我們可以天真地（沒有 vmap）實現(xiàn)它，如下所示：
def?naively_batched_linear(X_batched):
?return?jnp.stack([linear(x)?for?x?in?X_batched])
相反，通過使用 vmap 對線性進行向量化，我們可以一次性計算整個批次：文章來源：http://www.zghlxwxcb.cn/news/detail-770333.html
def?vmap_batched_linear(X_batched):
?return?vmap(linear)(X_batched)
本文由 mdnice 多平臺發(fā)布文章來源地址http://www.zghlxwxcb.cn/news/detail-770333.html
到了這里，關(guān)于JAX: 快如 PyTorch，簡單如 NumPy - 深度學(xué)習(xí)與數(shù)據(jù)科學(xué)的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
