JAX 陣列操作與自動微分技術核心概念

– 8 分鐘閱讀

JAX 作為一個高效能數值計算函式庫，其核心概念之一就是陣列操作。不同於 NumPy 的可變陣列，JAX 採用不可變陣列設計，更有利於平行計算和函式式程式設計。JAX 提供了 jax.numpy 高階介面和 jax.lax 低階介面，方便開發者根據需求選擇使用。型別系統的嚴格性有助於 JAX 進行更好的最佳化和錯誤檢查，而控制流原始碼和型別提升機制則賦予了 JAX 更強大的計算能力。在機器學習中，梯度計算至關重要，JAX 提供了手動、符號、數值和自動微分等多種計算梯度的方式，其中自動微分是深度學習的根本。

陣列在JAX

JAX提供了很多陣列操作的功能，例如可以讓我們建立、操作和轉換陣列。

切換到JAX NumPy-like API

JAX提供了一個NumPy-like API，可以讓我們更容易地使用JAX。

什麼是Array？

Array是JAX中的一個重要概念，它代表了一個多維度的陣列。

裝置相關操作

JAX提供了很多裝置相關操作的功能，例如可以讓我們在不同的裝置上執行陣列操作。

非同步派遣

非同步派遣是一種重要的技術，可以讓我們在執行陣列操作時不阻塞其他執行緒。

3.3 NumPy 與 JAX 的差異

在深入探討 JAX 之前，瞭解它與 NumPy 的差異是非常重要的。雖然 JAX 的設計初衷是要與 NumPy 相容，但它也引入了一些不同之處，以便更好地支援現代機器學習和高效能運算的需求。

3.3.1 不可變性（Immutability）

JAX 中的陣列是不可變的（immutable），這意味著一旦陣列被建立，它就不能被修改。這與 NumPy 中的陣列不同，NumPy 的陣列是可變的（mutable）。這種設計選擇使得 JAX 能夠更好地支援平行計算和函式語言程式設計。

3.3.2 型別（Types）

JAX 的型別系統比 NumPy 更加嚴格。JAX 需要明確指定陣列的型別，這使得 JAX 能夠進行更好的最佳化和錯誤檢查。

3.4 高階和低階介面：jax.numpy 和 jax.lax

JAX 提供了兩種不同的介面：高階介面 jax.numpy 和低階介面 jax.lax。高階介面提供了一種更簡單、更直觀的方式來使用 JAX，而低階介面提供了一種更靈活、更可定製的方式來控制 JAX 的行為。

3.4.1 控制流原始碼（Control Flow Primitives）

JAX 的控制流原始碼允許使用者定義自己的控制流程，這使得 JAX 能夠支援更複雜的計算。

3.4.2 型別提升（Type Promotion）

JAX 的型別提升機制允許使用者將陣列從一個型別提升到另一個型別，這使得 JAX 能夠支援更廣泛的計算。

4 計算梯度

計算梯度是機器學習中的一個基本問題。JAX 提供了多種方式來計算梯度，包括手動微分、符號微分、數值微分和自動微分。

4.1 不同方式的導數

JAX 支援多種方式來計算導數，包括：

  • 手動微分（Manual Differentiation）：使用者自己實作導數的計算。
  • 符號微分（Symbolic Differentiation）：使用符號運算來計算導數。
  • 數值微分（Numerical Differentiation）：使用數值方法來近似導數。
  • 自動微分（Automatic Differentiation）：JAX 自動計算導數。

內容解密：

上述內容介紹了 JAX 的基本概念和特性，包括其與 NumPy 的差異、不可變性、型別系統、控制流原始碼和型別提升機制。同時，也介紹了 JAX 中計算梯度的不同方式，包括手動微分、符號微分、數值微分和自動微分。這些內容為使用 JAX 進行高效能運算和機器學習提供了基礎知識。

import jax
import jax.numpy as jnp

# 定義一個函式
def my_function(x):
    return jnp.sin(x)

# 使用自動微分計算導數
grad_my_function = jax.grad(my_function)

# 測試導數
x = 1.0
print(grad_my_function(x))

圖表翻譯：

下面的 Plantuml 圖表展示了 JAX 中計算梯度的過程： 這個圖表展示了使用者可以選擇不同的方式來計算梯度，並且最終傳回計算出的導數。

4.2 自動微分計算梯度

在深度學習中，梯度計算是反向傳播演算法的核心。自動微分（autodiff）是一種高效計算梯度的方法。下面，我們將介紹如何使用 TensorFlow、PyTorch 和 JAX 進行梯度計算。

4.2.1 使用 TensorFlow 計算梯度

TensorFlow 提供了 tf.GradientTape 來計算梯度。以下是使用 tf.GradientTape 計算梯度的例子：

import tensorflow as tf

# 定義一個函式
def func(x):
    return x**2

# 建立一個 GradientTape 物件
with tf.GradientTape() as tape:
    # 記錄函式的輸入和輸出
    x = tf.Variable(2.0)
    y = func(x)

# 計算梯度
grad = tape.gradient(y, x)

print(grad)  # Output: 4.0

在這個例子中，我們定義了一個函式 func(x) = x**2，然後使用 tf.GradientTape 記錄函式的輸入和輸出。最後，我們使用 tape.gradient 計算梯度。

4.2.2 使用 PyTorch 計算梯度

PyTorch 提供了 torch.autograd 來計算梯度。以下是使用 torch.autograd 計算梯度的例子：

import torch

# 定義一個函式
def func(x):
    return x**2

# 建立一個 tensor 物件
x = torch.tensor(2.0, requires_grad=True)

# 計算梯度
y = func(x)
grad = torch.autograd.grad(y, x)

print(grad)  # Output: tensor(4.)

在這個例子中，我們定義了一個函式 func(x) = x**2，然後使用 torch.tensor 建立一個 tensor 物件。最後，我們使用 torch.autograd.grad 計算梯度。

4.2.3 使用 JAX 計算梯度

JAX 提供了 jax.grad 來計算梯度。以下是使用 jax.grad 計算梯度的例子：

import jax

# 定義一個函式
def func(x):
    return x**2

# 建立一個 array 物件
x = jax.numpy.array(2.0)

# 計算梯度
grad = jax.grad(func)(x)

print(grad)  # Output: 4.0

在這個例子中，我們定義了一個函式 func(x) = x**2，然後使用 jax.numpy.array 建立一個 array 物件。最後，我們使用 jax.grad 計算梯度。

4.3 前向和反向自動微分

自動微分可以分為前向自動微分和反向自動微分兩種。前向自動微分是指從輸入開始，逐步計算梯度；而反向自動微分是指從輸出開始，逐步計算梯度。

4.3.1 前向自動微分

前向自動微分是指從輸入開始，逐步計算梯度。以下是使用前向自動微分計算梯度的例子：

import jax

# 定義一個函式
def func(x):
    return x**2

# 建立一個 array 物件
x = jax.numpy.array(2.0)

# 計算梯度
grad = jax.jvp(func, (x,), (1.0,))

print(grad)  # Output: 4.0

在這個例子中，我們定義了一個函式 func(x) = x**2，然後使用 jax.jvp 計算梯度。

4.3.2 反向自動微分

反向自動微分是指從輸出開始，逐步計算梯度。以下是使用反向自動微分計算梯度的例子：

import jax

# 定義一個函式
def func(x):
    return x**2

# 建立一個 array 物件
x = jax.numpy.array(2.0)

# 計算梯度
grad = jax.vjp(func, x)

print(grad)  # Output: 4.0

在這個例子中，我們定義了一個函式 func(x) = x**2，然後使用 jax.vjp 計算梯度。

圖表翻譯：

@startuml
skinparam backgroundColor #FEFEFE
skinparam componentStyle rectangle

title JAX 陣列操作與自動微分技術核心概念

package "NumPy 陣列操作" {
    package "陣列建立" {
        component [ndarray] as arr
        component [zeros/ones] as init
        component [arange/linspace] as range
    }

    package "陣列操作" {
        component [索引切片] as slice
        component [形狀變換 reshape] as reshape
        component [堆疊 stack/concat] as stack
        component [廣播 broadcasting] as broadcast
    }

    package "數學運算" {
        component [元素運算] as element
        component [矩陣運算] as matrix
        component [統計函數] as stats
        component [線性代數] as linalg
    }
}

arr --> slice : 存取元素
arr --> reshape : 改變形狀
arr --> broadcast : 自動擴展
arr --> element : +, -, *, /
arr --> matrix : dot, matmul
arr --> stats : mean, std, sum
arr --> linalg : inv, eig, svd

note right of broadcast
  不同形狀陣列
  自動對齊運算
end note

@enduml

這個圖表展示了前向和反向自動微分的過程。從輸入開始，前向自動微分逐步計算梯度；而從輸出開始，反向自動微分逐步計算梯度。

內容解密：

在這個章節中，我們介紹瞭如何使用 TensorFlow、PyTorch 和 JAX 進行梯度計算。我們還介紹了前向和反向自動微分的概念和實作方式。透過這個章節，讀者可以瞭解如何使用自動微分計算梯度，並且可以應用於深度學習中。

5.2 即時編譯（JIT）內部機制

即時編譯（JIT）是一種在程式執行期間將中間碼轉換為機器碼的技術。要了解JIT的內部機制，我們需要探討JAX程式的中間表示形式，即Jaxpr。

Jaxpr：JAX程式的中間表示

Jaxpr是一種用於表示JAX程式的中間表示形式。它提供了一種平臺無關的方式來表示計算圖，使得JAX可以在不同硬體平臺上執行。Jaxpr的設計目的是為了最佳化計算圖的執行效率，同時也提供了一種靈活的方式來擴充套件JAX的功能。

XLA：加速線性代數運算

XLA（Accelerated Linear Algebra）是一種用於加速線性代數運算的技術。它提供了一種高效的方式來執行矩陣運算，同時也支援向量化運算。XLA是JAX的一個重要組成部分，它使得JAX可以在不同硬體平臺上高效地執行。

使用AOT編譯

AOT（Ahead-Of-Time）編譯是一種在程式編譯期間將程式碼轉換為機器碼的技術。使用AOT編譯可以提高程式的執行效率，因為它避免了在執行期間進行即時編譯。然而，AOT編譯也有一些限制，例如它需要事先知道程式的執行環境。

5.3 即時編譯的限制

即時編譯有一些限制，包括：

  • 純函式和不純函式：即時編譯只能對純函式進行最佳化，如果函式有副作用，則可能會導致錯誤的結果。
  • 精確數值：即時編譯可能會導致精確數值的損失，特別是在浮點數運算中。
  • 條件編譯：即時編譯可能會導致條件編譯失敗，特別是在複雜的條件陳述式中。
  • 編譯速度：即時編譯可能會導致編譯速度變慢，特別是在大型程式中。

類別方法和簡單函式

類別方法和簡單函式是兩種不同的函式型別。類別方法是指定義在類別中的函式，而簡單函式是指不依賴於類別的函式。即時編譯對這兩種函式型別有不同的最佳化策略。

練習5.1

練習5.1要求讀者實作一個簡單的即時編譯器，該編譯器可以將中間碼轉換為機器碼。

從效能最佳化視角來看，JAX提供的陣列操作和即時編譯（JIT）功能對於高效能數值計算至關重要。深入剖析JAX的Jaxpr中間表示和XLA的整合，可以發現JAX如何透過硬體加速和編譯最佳化提升運算效率。分析Jaxpr如何將JAX程式碼轉換成平臺無關的表示形式，以及XLA如何針對特定硬體進行最佳化，可以理解JAX高效能的關鍵。然而，JIT的限制，例如純函式要求、潛在的數值精確度損失以及條件編譯的挑戰，需要開發者仔細考量。對於追求極致效能的應用，AOT編譯提供了一個替代方案，但需要權衡其佈署的靈活性。對於重視效能的開發者，深入理解JAX的運作機制，並根據應用場景選擇合適的編譯策略，才能最大限度地發揮JAX的效能優勢。玄貓認為，JAX在高效能數值計算和機器學習領域的應用前景廣闊，值得深入研究和應用。

玄貓 BlackCat

玄貓 BlackCat

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。