jax.numpy: NumPy와 닮은 API

jax.numpy의 약속은 단순해. NumPy의 API를 거의 그대로 제공하되 연산은 XLA 위에서 실행한다는 거야. 그래서 NumPy 코드에서 JAX로 옮기는 첫 단계는 import 한 줄을 바꾸는 것부터 시작해.

import numpy as np
import jax.numpy as jnp

a = np.array([1, 2, 3])
b = jnp.array([1, 2, 3])

print(np.sum(a))   # 6
print(jnp.sum(b))  # 6

# 거의 모든 게 그대로 — 함수 이름, 시그니처, 동작
np.zeros((3, 3));   jnp.zeros((3, 3))
np.linspace(0,1,5); jnp.linspace(0,1,5)
np.dot(a, a);       jnp.dot(b, b)

함수 이름과 시그니처는 비슷하지만 실행 모델에는 중요한 차이가 있어.

불변 배열: jnp.array는 직접 변경할 수 없어. a[0] = 5 대신 a = a.at[0].set(5)를 사용해.
자동 장치 배치: 배열을 만들면 사용 가능한 가속기에 배치된다.
float32 기본 dtype: NumPy의 기본값인 float64와 다르다.
별도의 난수 API: jax.random과 명시적인 키를 사용한다. Track 8에서 자세히 배운다.

⚠️ 첫 번째 함정

NumPy 코드를 그대로 옮기면 인덱스 할당에서 거의 틀림없이 막혀. .at[].set()을 이용한 함수형 갱신 패턴을 익혀야 해.

Code

import numpy as np
import jax.numpy as jnp

# Creating arrays — identical syntax
np_arr = np.array([1.0, 2.0, 3.0])
jnp_arr = jnp.array([1.0, 2.0, 3.0])

# Zeros, ones, ranges — identical
np_zeros = np.zeros((3, 4))
jnp_zeros = jnp.zeros((3, 4))

np_range = np.arange(0, 10, 2)
jnp_range = jnp.arange(0, 10, 2)

np_lin = np.linspace(0, 1, 50)
jnp_lin = jnp.linspace(0, 1, 50)

# Math — identical
print(np.sum(np_arr ** 2))     # 14.0
print(jnp.sum(jnp_arr ** 2))  # 14.0

import jax.numpy as jnp

# From Python data
a = jnp.array([[1, 2], [3, 4]])

# Filled arrays
zeros = jnp.zeros((3, 4))           # 3x4 of zeros
ones = jnp.ones((2, 3))             # 2x3 of ones
full = jnp.full((2, 2), fill_value=7.0)  # 2x2 of 7.0

# Ranges
r = jnp.arange(0, 10, 0.5)          # [0.0, 0.5, 1.0, ..., 9.5]
l = jnp.linspace(0, 1, 5)           # [0.0, 0.25, 0.5, 0.75, 1.0]

# Identity and diagonal
eye = jnp.eye(3)                     # 3x3 identity matrix
diag = jnp.diag(jnp.array([1, 2, 3]))  # 3x3 diagonal matrix

# Random arrays (JAX uses explicit PRNG keys, covered in Track 3)
import jax
key = jax.random.PRNGKey(0)
rand = jax.random.normal(key, shape=(3, 4))  # 3x4 standard normal

x = jnp.array([[1.0, 2.0, 3.0],
                [4.0, 5.0, 6.0]])

print(x.shape)   # (2, 3)
print(x.dtype)   # float32
print(x.ndim)    # 2
print(x.size)    # 6

Exercise

전에 작성한 30줄 안팎의 NumPy 스크립트를 하나 골라. import numpy as np를 import jax.numpy as jnp로, np.를 jnp.로 바꾼 뒤 실행해. 어떤 호출이 오류를 내고 어떤 동작이 조용히 달라지는지 정확히 기록해. 아직 고치지 마. 이번 연습의 목적은 차이를 진단하는 거야.

⚠️ 첫 번째 함정

Code

External links

Exercise

Progress

댓글 0