🔪JAX - 세부적인 특징들 (JAX - The Sharp Bits)🔪

저자 : Anselm Levskaya & Mattew Johnson

역자 : 장진우(wdfrty1234@gmail.com)

검수 : 이영빈, 박정현

이탈리아 변두리를 걷다보면, 사람들이 JAX에 대해 *“una anima di pura programmazione funzionale(순수한 함수형 프로그래밍의 영혼)”*이라는 용어로 묘사하는 사람들이 많다는 것을 알 수 있습니다.

JAX는 수치해석 프로그램의 변환을 위한 언어로, CPU 혹은 가속기(GPU/TPU)에서 수치형 프로그램을 컴파일하여 동작시킬 수 있습니다. 특정한 제약조건을 충족하는 경우, JAX는 수치 및 과학 프로그램에서 잘 동작합니다!

import numpy as np
from jax import grad, jit
from jax import lax
from jax import random
import jax
import jax.numpy as jnp

🔪순수 함수 (Pure functions)

JAX는 (모든 입력 데이터가 함수의 매개변수를 통해 전달되고, 모든 출력 결과가 함수의 결과를 통해서 나오는) 순수 함수에서 잘 동작하도록 설계되어 있습니다.

(역자 주) 순수 함수는 같은 입력이 주어진다면 항상 같은 결과를 반환하는 함수를 의미합니다.

하지만, 하지만 순수함수가 아닌 경우 JAX는 Python 인터프리터와 다르게 동작할 수 있습니다. 아래는 그 예시 입니다. 이와 같은 예는 JAX 시스템에서의 동작이 보장되지 않습니다. 따라서, JAX를 사용하는 가장 적절한 방법은 함수적으로 순수한 Python함수에 대해서만 사용하는 것입니다.

def impure_print_side_effect(x):
  print("Executing function")  # This is a side-effect
  return x

# The side-effects appear during the first run
print ("First call: ", jit(impure_print_side_effect)(4.))

# Subsequent runs with parameters of same type and shape may not show the side-effect
# This is because JAX now invokes a cached compilation of the function
print ("Second call: ", jit(impure_print_side_effect)(5.))

# JAX re-runs the Python function when the type or shape of the argument changes
print ("Third call, different type: ", jit(impure_print_side_effect)(jnp.array([5.])))

g = 0.
def impure_uses_globals(x):
  return x + g

# JAX captures the value of the global during the first run
print ("First call: ", jit(impure_uses_globals)(4.))
g = 10.  # Update the global

# Subsequent runs may silently use the cached value of the globals
print ("Second call: ", jit(impure_uses_globals)(5.))

# JAX re-runs the Python function when the type or shape of the argument changes
# This will end up reading the latest value of the global
print ("Third call, different type: ", jit(impure_uses_globals)(jnp.array([4.])))

g = 0.
def impure_saves_global(x):
  global g
  g = x
  return x

# JAX runs once the transformed function with special Traced values for arguments
print ("First call: ", jit(impure_saves_global)(4.))
print ("Saved global: ", g)  # Saved global has an internal JAX value

만약 Python 함수가 실제로 스테이트풀 객체를 내부적으로 사용하더라도, 이를 외부에서 읽거나 쓰지 않는다면 함수적으로 순수하다고 볼 수 있습니다.

def pure_uses_internal_state(x):
  state = dict(even=0, odd=0)
  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을 사용하려는 JAX 함수나 어떤 제어 흐름 구성요소에서 iterators를 사용하는 것은 권장되지 않습니다. 그 이유는 iterator(반복자)가 다음 요소를 가져오기 위한 상태(state)를 찾기 위한 파이썬 객체이기 때문입니다. 이러한 이유로, iterators는 JAX의 함수적 프로그래밍 모델과 호환되지 않습니다. 아래 코드는 JAX에서 iterators를 사용하려는 부적절한 시도들에 대한 예시를 보여줍니다. 이러한 예제들 대부분은 오류를 반환하지만, 어떤 경우에는 예상치 못한 결과를 보여줄 수도 있습니다.

import jax.numpy as jnp
import jax.lax as lax
from jax import make_jaxpr

# lax.fori_loop
array = jnp.arange(10)
print(lax.fori_loop(0, 10, lambda i,x: x+array[i], 0)) # expected result 45
iterator = iter(range(10))
print(lax.fori_loop(0, 10, lambda i,x: x+next(iterator), 0)) # unexpected result 0

# lax.scan
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.) # throws error

# lax.cond
array_operand = jnp.array([0.])
lax.cond(True, lambda x: x+1, 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) # throws error

🔪In-Place 업데이트 (In-Place Updates)

Numpy를 사용할 때, 여러분들은 종종 이렇게 사용할 것입니다.

numpy_array = np.zeros((3,3), dtype=np.float32)
print("original array:")
print(numpy_array)

# In place, mutating update
numpy_array[1, :] = 1.0
print("updated array:")
print(numpy_array)

하지만, 만약 JAX device array에 in-place로 업데이트를 시도하면, 오류가 발생하는 것을 볼 수 있습니다. (☉_☉)

%xmode Minimal

jax_array = jnp.zeros((3,3), dtype=jnp.float32)

# In place update of JAX's array will yield an error!
jax_array[1, :] = 1.0

변수의 in-place 변형을 허용하는 것은 프로그램의 분석과 변환을 어렵게 만드는 요인입니다. JAX에서는 프로그램이 순수 함수여야 한다는 것을 기억합시다.

JAX에서는 in-place 기법 대신에, JAX에서는 JAX array에 .at 속성을 이용하여 함수적 배열 업데이트를 수행합니다. (.at property on JAX arrays.)

⚠️ jit된 코드 와 lax.while_loop 또는 lax.fori_loop 내부에서 슬라이스의 크기는 인수 값의 함수가 아니라 인수 형태의 함수여야 가능합니다. - 슬라이스 시작 인덱스에는 그러한 제한이 없습니다. 아래의 제어 흐름 부분에서 이러한 제약에 대한 정보를 확인할 수 있습니다.

배열 업데이트 : x.at[idx].set(y)

예를 들어, 업데이트는 아래와 같이 작성할 수 있습니다.

updated_array = jax_array.at[1, :].set(1.0)
print("updated array:\n", updated_array)

JAX의 업데이트 함수는 NumPy와는 다르게 out-of-place로 동작합니다. 즉, 업데이트된 배열은 새 배열로 반환되며 원래 배열은 업데이트로 수정되지 않습니다.

print("original array unchanged:\n", jax_array)

하지만, jit으로 컴파일 된 코드 내에서 x.at[idx].set(y)의 입력 값 x가 재사용되지 않는 경우, 컴파일러는 in-place로 배열이 업데이트 되도록 최적화할 수 있습니다.

다른 연산과 함께 배열 업데이트

인덱스가 지정된 배열의 업데이트는 단순히 값을 덮어쓰는 것에만 제한되지는 않습니다. 예를 들어, 아래의 예시와 같이 인덱스에 덧셈을 하는 연산도 수행할 수 있습니다.

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)

보다 더 자세한 인덱스된 배열의 업데이트에 관련하여서는, 해당 문서를 참고해주세요. (documentatiaon for the .at property)

🔪 범위를 벗어난 인덱싱 (Out-of-Bound Indexing)

NumPy에서는 여러분이 인덱스 배열의 인덱스 범위를 벗어나는 동작을 수행하면 아래와 같은 에러를 볼 수 있습니다.

np.arange(10)[11]

하지만, 가속기에서 동작하는 코드로부터 에러를 발생시키는 것은 어렵거나 심지어는 불가능할 수 있습니다. 그러므로, JAX는 배열의 범위를 벗어나는 인덱싱에 대해서 오류가 아닌 동작을 선택해야 합니다. (유효하지 않은 부동 소수점의 산술적 결과가 NaN이 되는 것과 유사합니다.). 만약 인덱싱 작업이 배열 인덱스 업데이트(예: index_add 또는 scatter-유사한 기본 요소)인 경우, 범위를 벗어난 인덱스의 업데이트는 건너뜁니다. 작업이 배열 인덱스 검색(예: NumPy 인덱싱 또는 gather-유사 기본 요소)인 경우, 무언가를 반환해야 하므로 인덱스가 배열의 범위에 고정됩니다. 예를 들어, 아래의 인덱싱 동작에서는 배열의 마지막 값이 반환될 것입니다.

jnp.arange(10)[11]

인덱스 검색에 대한 이러한 동작으로 인해 jnp.nanargmin 및 jnp.nanargmax와 같은 함수는 NaN으로 구성된 슬라이스에 대해 -1을 반환하지만 Numpy는 오류를 발생시킵니다.

위에서 설명한 두 가지 동작이 서로 역의 관계가 아니기 때문에, 역방향 자동 미분(인덱스 업데이트를 인덱스 검색으로 변환하고 그 반대로 전환)은 범위를 벗어난 인덱싱의 의미를 보존하지 않습니다. 따라서 JAX의 범위를 벗어난 인덱싱을 정의되지 않은 동작으로 생각하는 것이 좋습니다.

🔪 비배열 입력: NumPy vs. Jax (Non-array inputs: NumPy vs. JAX)

NumPy는 일반적으로 Python의 리스트 또는 튜플을 API 함수에 대한 입력으로 사용합니다.

np.sum([1, 2, 3])

JAX는 이와 다르게 일반적으로 유용한 오류를 반환합니다.

jnp.sum([1, 2, 3])

이는 의도적으로 설계된 결과입니다. 그 이유는 추적된 함수에 리스트나 튜플을 전달하게 되면 감지하기 어려운 조용한 성능의 저하가 발생할 수 있기 때문입니다.

예를 들어, 리스트의 입력을 허용하는 다음 버전의 jnp.sum을 고려해봅시다.

def permissive_sum(x):
  return jnp.sum(jnp.array(x))

x = list(range(10))
permissive_sum(x)

결과는 예상한 대로이지만 여기에는 잠재적인 성능 문제가 숨겨져 있습니다. JAX의 추적 및 JIT 컴파일 모델에서 Python의 리스트 혹은 튜플의 각 요소가 별도의 JAX 변수로 취급되며, 이러한 변수들은 개별적으로 처리되어 디바이스로 전송됩니다. 이는 위의 permissive_sum 함수에 대한 jaxpr에서 볼 수 있습니다.

make_jaxpr(permissive_sum)(x)

리스트의 각 항목은 별도의 입력으로 처리되므로, 리스트의 크기에 따라 선형적으로 증가하는 추적 및 컴파일 오버헤드가 발생합니다. 이러한 문제를 방지하기 위해, JAX는 리스트 및 튜플을 배열로 임시적으로 변환하는 것을 피합니다.

따라서, 튜플 또는 리스트를 JAX 함수에 전달하려면 먼저 명시적으로 배열로 변환한 후 전달해야 합니다.

jnp.sum(jnp.array(x))

🔪 난수 (Random Numbers)

rand()로 인해 결과가 의심스러운 모든 과학 논문들이 도서관 책장에서 사라진다면 각 책장에는 주먹만한 간격이 생길 겁니다. - Numerical Recipes

RNGs와 State

여러분들은 NumPy 및 기타 라이브러리의 스테이트풀 의사 난수 생성기(PRNG)에 익숙할 것입니다. 이러한 라이브러리들은 의사 난수의 소스를 제공하기 위해 많은 세부 정보들을 백그라운드에서 유용하게 숨깁니다.

print(np.random.random())
print(np.random.random())
print(np.random.random())

백그라운드에서 numpy는 Mersenne Twister PRNG를 사용하여 의사 난수 기능을 강화합니다. PRNG의 주기는 \(2^{19937} - 1\)이고 어느 시점에서든 624개의 32비트 부호 없는 정수와 이 “엔트로피”가 얼마나 많이 사용되었는지에 대한 위치로 설명할 수 있습니다.

np.random.seed(0)
rng_state = np.random.get_state()
# print(rng_state)
# --> ('MT19937', array([0, 1, 1812433255, 1900727105, 1208447044,
#       2481403966, 4042607538,  337614300, ... 614 more numbers...,
#       3048484911, 1796872496], dtype=uint32), 624, 0, 0.0)

이 의사 난수 상태 벡터는 난수가 필요할 때마다 백그라운드에서 자동적으로 업데이트되어 Mersenne twister 상태 벡터의 uint32 중 2개를 “소비”합니다.

_ = np.random.uniform()
rng_state = np.random.get_state()
#print(rng_state)
# --> ('MT19937', array([2443250962, 1093594115, 1878467924,
#       ..., 2648828502, 1678096082], dtype=uint32), 2, 0, 0.0)

# Let's exhaust the entropy in this PRNG statevector
for i in range(311):
  _ = np.random.uniform()
rng_state = np.random.get_state()
#print(rng_state)
# --> ('MT19937', array([2443250962, 1093594115, 1878467924,
#       ..., 2648828502, 1678096082], dtype=uint32), 624, 0, 0.0)

# Next call iterates the RNG state for a new batch of fake "entropy".
_ = np.random.uniform()
rng_state = np.random.get_state()
# print(rng_state)
# --> ('MT19937', array([1499117434, 2949980591, 2242547484,
#      4162027047, 3277342478], dtype=uint32), 2, 0, 0.0)

Magic PRNG 상태의 문제는 서로 다른 스레드, 프로세스 및 장치에서 사용 및 업데이트 되는 방식에 대해 추론하기 어렵고 엔트로피 생성 및 소비에 대한 세부 정보가 최종 사용자에게 숨겨져 있을 때 문제를 일으키기 매우 쉽다는 것입니다.

또한, Mersenne Twister PRNG는 많은 문제가 있는 것으로 알려져 있으며, 2.5Kb의 큰 상태 크기를 가지고 있어 초기화 문제를 야기할 수 있습니다. 또한, 최신 BigCrush 테스트를 만족하지 못하고 일반적으로 느리다는 단점이 있습니다.

JAX PRNG

JAX는 대신 PRNG 상태를 명시적으로 전달하고 반복하여 엔트로피 생성 및 소비를 처리하는 명시적 PRNG를 구현했습니다. JAX는 분할 가능한 최신 Threefry counter 기반 PRNG를 사용합니다(Threefry counter-based PRNG). 즉, 이러한 설계를 통해 PRNG 상태를 병렬 확률적 생성을 위해 사용하기 위해 새로운 PRNG로 분기할 수 있습니다.

무작위 상태는 키(key)라고 부르는 두 개의 unsigned-int32로 설명됩니다.

from jax import random
key = random.PRNGKey(0)
key

JAX의 임의 함수는 PRNG 상태에서 의사 난수를 생성하지만 상태를 변경하지는 않습니다!

동일한 상태를 재사용하는 행위는 슬픔과 단조로움을 유발하며 결국 최종 사용자에게 혼돈을 불어넣는 결과를 초래할 수 있습니다!
(Reusing the same state will cause sadness and monotony, depriving the end user of lifegiving chaos:)

print(random.normal(key, shape=(1,)))
print(key)
# No no no!
print(random.normal(key, shape=(1,)))
print(key)

대신, 새로운 의사 난수가 필요할 때마다 PRNG를 분할하여 사용 가능한 하위 키를 얻습니다.

print("old key", key)
key, subkey = random.split(key)
normal_pseudorandom = random.normal(subkey, shape=(1,))
print("    \---SPLIT --> new key   ", key)
print("             \--> new subkey", subkey, "--> normal", normal_pseudorandom)

새로운 난수가 필요할 때마다 키를 전파하고 새 하위 키를 만듭니다.

print("old key", key)
key, subkey = random.split(key)
normal_pseudorandom = random.normal(subkey, shape=(1,))
print("    \---SPLIT --> new key   ", key)
print("             \--> new subkey", subkey, "--> normal", normal_pseudorandom)

한 번에 둘 이상의 하위키를 만들 수 있습니다.

key, *subkeys = random.split(key, 4)
for subkey in subkeys:
  print(random.normal(subkey, shape=(1,)))

🔪 제어 흐름 (Control Flow)

✔ python control_flow + autodiff ✔

Python 함수에 grad를 적용하려는 경우 Autograd(또는 Pytorch 또는 TF Eager)를 사용하는 것처럼 문제 없이 일반적인 Python 제어 흐름 구성을 사용할 수 있습니다.

def f(x):
  if x < 3:
    return 3. * x ** 2
  else:
    return -4 * x

print(grad(f)(2.))  # ok!
print(grad(f)(4.))  # ok!

Python control flow + JIT

jit와 함께 제어 흐름을 사용하는 것은 더 복잡하며 기본적으로 더 많은 제약이 있습니다.

이 예시는 동작합니다.

@jit
def f(x):
  for i in range(3):
    x = 2 * x
  return x

print(f(3))

아래 예시도 동작합니다.

@jit
def g(x):
  y = 0.
  for i in range(x.shape[0]):
    y = y + x[i]
  return y

print(g(jnp.array([1., 2., 3.])))

하지만 이 예시는 기본적으로 동작하지 않습니다.

@jit
def f(x):
  if x < 3:
    return 3. * x ** 2
  else:
    return -4 * x

# This will fail!
f(2)

왜 그럴까!?

함수를 jit 컴파일할 때 일반적으로 컴파일된 코드를 캐시하여 다양한 인수 값에 대해 작동하는 함수 버전을 컴파일하여 재사용할 수 있도록 합니다. 이러한 방식으로 각 함수 평가마다 다시 컴파일할 필요가 없습니다.

예를 들어 jnp.array([1., 2., 3.], jnp.float32) 배열에서 @jit 함수를 평가하기 위해 jnp.array([4., 5., 6.], jnp.float32)에서 사용했던 컴파일된 코드를 재사용하여 하여 컴파일에 수행되는 시간을 절약할 수 있습니다.

JAX는 Python 코드의 다양한 인수 값에 유효한 뷰를 얻기 위해 JAX는 가능한 입력 집합을 나타내는 추상 값으로 코드를 추적합니다. 다양한 추상화 수준이 있으며 서로 다른 변환은 서로 다른 추상화 수준을 사용합니다.

기본적으로 jit은 ShapedArray 추상화 수준에서 코드를 추적합니다. 각 추상 값은 고정된 모양과 dtype을 가지는 모든 배열 값의 집합을 나타냅니다. 예를 들어 추상 값 ShapedAray((3,), jnp.float32)를 사용하여 추적하면 해당 배열 세트의 구체적인 값에 대해 재사용할 수 있는 함수의 뷰를 얻을 수 있습니다. 즉, 컴파일 시간을 줄일 수 있다는 것을 의미합니다.

그러나 여기에는 장단점이 있습니다. 특정 구체적인 값이 결정되지 않은 ShapedArray((), jnp.float32)에서 Python 함수를 추적하는 경우 if x < 3과 같은 줄에 도달하면 표현식 x < 3은 {True, False} 집합을 나타내는 추상 ShapedArray((), jnp.bool_)로 평가됩니다. Python이 이를 구체적인 True 또는 False로 강제하려고 하면 오류가 발생합니다. 어떤 분기를 선택해야 할지 모르고 추적을 계속할 수 없습니다! 단점은 추상화 수준이 높을수록 Python 코드에 대한 보다 일반적인 뷰를 얻을 수 있지만(따라서 재컴파일을 줄일 수 있습니다.) 추적을 완료하려면 Python 코드에 더 많은 제약이 필요하다는 것입니다.

좋은 소식은, 이 트레이드오프를 직접 제어할 수 있다는 것입니다. 보다 정밀한 추상 값에 대한 jit 추적을 통해 추적 가능성 제약을 완화할 수 있습니다. 예를 들어 jit에 static_argnums 인수를 사용하여 일부 인수의 구체적인 값을 추적하도록 지정할 수 있습니다. 다음은 이에 대한 예제 함수입니다.

def f(x):
  if x < 3:
    return 3. * x ** 2
  else:
    return -4 * x

f = jit(f, static_argnums=(0,))

print(f(2.))

루프를 포함한 또다른 예제입니다.

def f(x, n):
  y = 0.
  for i in range(n):
    y = y + x[i]
  return y

f = jit(f, static_argnums=(1,))

f(jnp.array([2., 3., 4.]), 2)

static_argnums를 이용한 효과로 인해 루프가 정적으로 펼쳐집니다. JAX는 또한 Unshaped와 같은 더 높은 수준의 추상화에서 추적할 수 있지만 현재 변환의 기본값은 아닙니다.

⚠️ 인수 값에 따라 형태가 바뀌는 함수

이러한 제어 흐름 문제는 보다 미묘한 방식으로도 나타납니다. jit 하려는 수치 함수는 내부 배열의 모양을 인수 값에 따라 특정할 수 없습니다(인수 모양에 따라 특정하는 것은 괜찮습니다). 간단한 예로 입력 변수 길이에 따라 출력이 달라지는 함수를 만들어 보겠습니다.

def example_fun(length, val):
  return jnp.ones((length,)) * val
# un-jit'd works fine
print(example_fun(5, 4))

bad_example_jit = jit(example_fun)
# this will fail:
bad_example_jit(10, 4)

# static_argnums tells JAX to recompile on changes at these argument positions:
good_example_jit = jit(example_fun, static_argnums=(0,))
# first compile
print(good_example_jit(10, 4))
# recompiles
print(good_example_jit(5, 4))

static_argnums는 예제에서 길이의 변경이 잦지 않은 경우에는 편리할 수 있지만 변경이 잦은 경우 재앙이 될 수 있습니다!

마지막으로 함수에 전역적인 부수효과들이 있는 경우 JAX의 추적 프로그램으로 인해 이상한 일이 발생할 수 있습니다. 일반적인 문제는 jit’d 함수 내에서 배열을 출력하려고 시도할 때 발생할 수 있습니다.

@jit
def f(x):
  print(x)
  y = 2 * x
  print(y)
  return y
f(2)

구조적 제어 흐름 기본 요소

JAX에는 제어 흐름에 대한 다양한 옵션이 많이 있습니다. 예를 들어 재컴파일을 피하고 추적 가능한 제어 흐름을 사용하면서 큰 루프를 풀고 싶지 않다면 아래의 4가지 구조적 제어 흐름 기본 구조를 사용할 수 있습니다.

• lax.cond differentiable

• lax.while_loop fwd-mode-differentiable

• lax.fori_loop fwd-mode-differentiable in general; fwd and rev-mode differentiable if endpoints are static.

• lax.scan differentiable

cond

def cond(pred, true_fun, false_fun, operand):
  if pred:
    return true_fun(operand)
  else:
    return false_fun(operand)

from jax import lax

operand = jnp.array([0.])
lax.cond(True, lambda x: x+1, lambda x: x-1, operand)
# --> array([1.], dtype=float32)
lax.cond(False, lambda x: x+1, lambda x: x-1, operand)
# --> array([-1.], dtype=float32)

jax.lax에는 동적 조건에 따라 분기할 수 있는 다른 두 개의 함수가 제공됩니다.

• lax.select는 lax.cond의 배치 버전이지만, 선택지는 이전에 계산된 배열로 표현됩니다.

• lax.switch는 lax.cond와 유사하지만, 어떤 수의 호출 가능한 선택지 사이에 전환할 수 있습니다.

또한, jax.numpy에서는 이러한 함수에 대한 다수의 Numpy 스타일 인터페이스가 제공됩니다.

• jnp.where는 3개의 인수가있는 lax.select의 Numpy 스타일 래퍼(wrapper)입니다.

• jnp.piecewise는 lax.switch의 Numpy 스타일 래퍼(wrapper)이지만, 단일 스칼라 인덱스 대신에 불리언 조건의 목록에 따라 전환합니다.

• jnp.select는 jnp.piecewise와 유사한 API를 가지지만, 선택지는 사전 계산된 배열로 제공됩니다. 결과적으로 lax.select의 여러 호출로 구현됩니다.

while_loop

def while_loop(cond_fun, body_fun, init_val):
  val = init_val
  while cond_fun(val):
    val = body_fun(val)
  return val

init_val = 0
cond_fun = lambda x: x<10
body_fun = lambda x: x+1
lax.while_loop(cond_fun, body_fun, init_val)
# --> array(10, dtype=int32)

fori_loop

def fori_loop(start, stop, body_fun, init_val):
  val = init_val
  for i in range(start, stop):
    val = body_fun(i, val)
  return val

init_val = 0
start = 0
stop = 10
body_fun = lambda i,x: x+i
lax.fori_loop(start, stop, body_fun, init_val)
# --> array(45, dtype=int32)

Summary

🔪 동적 형태 (Dynamic Shapes)

jax.jit, jax.vmap, jax.grad 등과 같은 변환 내에서 사용되는 JAX 코드는 모든 출력 배열과 중간 배열이 정적 모양을 가져야 합니다. 즉, 모양은 다른 배열 내의 값에 의존하지 않아야 합니다.

예를 들어, jnp.nansum의 버전을 직접 구현하려면 다음과 같이 시작할 수 있습니다

def nansum(x):
  mask = ~jnp.isnan(x)  # boolean mask selecting non-nan values
  x_without_nans = x[mask]
  return x_without_nans.sum()

JIT 및 기타 변환 외부에서는 예상대로 작동합니다.

x = jnp.array([1, 2, jnp.nan, 3, 4])
print(nansum(x))

jax.jit 또는 다른 변환을 이 함수에 적용하려고 하면 오류가 발생합니다.

jax.jit(nansum)(x)

문제는 x_without_nans의 크기가 x 내의 값에 따라 동적으로 결정되다는 점입니다. 종종 JAX에서는 다른 방법을 통해 동적으로 크기 조정된 배열의 필요성을 해결할 수 있습니다. 예를 들어 jnp.where의 3개 인수 형식을 사용하여 NaN 값을 0으로 대체함으로써 동적 모양을 피하면서도 동일한 결과를 계산할 수 있습니다.

@jax.jit
def nansum_2(x):
  mask = ~jnp.isnan(x)  # boolean mask selecting non-nan values
  return jnp.where(mask, x, 0).sum()

print(nansum_2(x))

동적 모양의 배열이 발생하는 다른 상황에서도 유사한 트릭을 사용할 수 있습니다.

🔪 NaNs

NaNs 디버깅

함수 또는 그래디언트에서 NaN이 발생하는 위치를 추적하려면 다음과 같이 NaN 검사기를 켤 수 있습니다.

• JAX_DEBUG_NANS=True 환경 변수 설정

• 메인 파일 상단에  from jax.config import config 와 config.update("jax_debug_nans",True) 를 추가하세요.

• 메인 파일에 from jax.config import config 와 config.parse_flags_with_absl() 를 추가하세요. 그런 다음 명령 줄 플래그에 -jax_debug_nans=True 을 이용하여 옵션을 설정하세요.

이로 인해 NaN 생성 즉시 계산 오류가 발생합니다. 이 옵션을 켜면 XLA에서 생성된 모든 부동 소수점 유형 값에 nan 검사가 추가됩니다. 즉, @jit에 속하지 않는 모든 기본 작업에 대해 값을 다시 호스트로 가져와 ndarry로 검사합니다. @jit 하위에 있는 코드의 경우 모든 @jit함수의 출력을 검사하고 NaN이 있는 경우 최적화되지 않은 op-by-op 모드에서 함수를 다시 실행하여 한 번에 한 레벨씩 @jit을 제거합니다.

@jit에서만 발생하지만 최적화되지 않은 모드에서는 생성되지 않는 nan과 같은 까다로운 상황이 발생할 수 있습니다. 이 경우 경고 메시지가 출력되지만 코드는 계속 실행됩니다.

그래디언트 평가의 역방향 패스에서 nans가 생성되는 경우 스택 추적에서 몇 프레임 위로 예외가 발생하면 backward_pass 함수 내부로 진입합니다. 이 함수는 기본적으로 기본 작업 시퀀스를 역순으로 수행하는 간단한 jaxpr 인터프리터입니다. 아래 예에서 env JAX_DEBUG_NANS=True ipython 명령줄을 사용하여 ipython repl을 시작한 후, 다음을 실행했습니다.

import jax.numpy as jnp

jnp.divide(0., 0.)

생성된 NaN이 잡혔습니다. %debug를 실행하면 사후 디버거를 얻을 수 있습니다. 이것은 아래 예제와 같이 @jit 으로 감싸진 함수에서도 작동합니다.

In [1]: import jax.numpy as jnp

In [2]: jnp.divide(0., 0.)
---------------------------------------------------------------------------
FloatingPointError                        Traceback (most recent call last)
<ipython-input-2-f2e2c413b437> in <module>()
----> 1 jnp.divide(0., 0.)

.../jax/jax/numpy/lax_numpy.pyc in divide(x1, x2)
    343     return floor_divide(x1, x2)
    344   else:
--> 345     return true_divide(x1, x2)
    346
    347

.../jax/jax/numpy/lax_numpy.pyc in true_divide(x1, x2)
    332   x1, x2 = _promote_shapes(x1, x2)
    333   return lax.div(lax.convert_element_type(x1, result_dtype),
--> 334                  lax.convert_element_type(x2, result_dtype))
    335
    336

.../jax/jax/lax.pyc in div(x, y)
    244 def div(x, y):
    245   r"""Elementwise division: :math:`x \over y`."""
--> 246   return div_p.bind(x, y)
    247
    248 def rem(x, y):

... stack trace ...

.../jax/jax/interpreters/xla.pyc in handle_result(device_buffer)
    103         py_val = device_buffer.to_py()
    104         if np.any(np.isnan(py_val)):
--> 105           raise FloatingPointError("invalid value")
    106         else:
    107           return DeviceArray(device_buffer, *result_shape)

FloatingPointError: invalid value

@jit
def f(x, y):
    a = x * y
    b = (x + y) / (x - y)
    c = a + 2
    return a + b * c

In [4]: from jax import jit

In [5]: @jit
   ...: def f(x, y):
   ...:     a = x * y
   ...:     b = (x + y) / (x - y)
   ...:     c = a + 2
   ...:     return a + b * c
   ...:

In [6]: x = jnp.array([2., 0.])

In [7]: y = jnp.array([3., 0.])

In [8]: f(x, y)
Invalid value encountered in the output of a jit function. Calling the de-optimized version.
---------------------------------------------------------------------------
FloatingPointError                        Traceback (most recent call last)
<ipython-input-8-811b7ddb3300> in <module>()
----> 1 f(x, y)

 ... stack trace ...

<ipython-input-5-619b39acbaac> in f(x, y)
      2 def f(x, y):
      3     a = x * y
----> 4     b = (x + y) / (x - y)
      5     c = a + 2
      6     return a + b * c

.../jax/jax/numpy/lax_numpy.pyc in divide(x1, x2)
    343     return floor_divide(x1, x2)
    344   else:
--> 345     return true_divide(x1, x2)
    346
    347

.../jax/jax/numpy/lax_numpy.pyc in true_divide(x1, x2)
    332   x1, x2 = _promote_shapes(x1, x2)
    333   return lax.div(lax.convert_element_type(x1, result_dtype),
--> 334                  lax.convert_element_type(x2, result_dtype))
    335
    336

.../jax/jax/lax.pyc in div(x, y)
    244 def div(x, y):
    245   r"""Elementwise division: :math:`x \over y`."""
--> 246   return div_p.bind(x, y)
    247
    248 def rem(x, y):

 ... stack trace ...

이 코드는 @jit 함수의 출력에서 nan을 발견하면 최적화되지 않은 코드를 호출하므로 여전히 명확한 스택을 추적할 수 있습니다. 그리고 %debug로 사후 디버거를 실행하여 오류를 파악하기 위해 모든 값을 검사할 수 있습니다.

⚠️ 디버깅하지 않는 경우 NaN 검사기를 켜서는 안 됩니다. 많은 장치-호스트 왕복 및 성능 저하가 발생할 수 있기 때문입니다!

⚠️ NaN 검사기는 pmap에서 작동하지 않습니다. pmap 코드에서 nans를 디버깅하려면 pmap을 vmap으로 교체해야 합니다.

🔪 Double (64bit) 정밀도 (Double (64bit) precision)

현재 JAX는 기본적으로 NumPy API가 피연산자를 강제로 더블형(double)으로 변환하는 경향을 완화하기 위해 단정밀도(single-precision) 숫자를 강제로 적용하고 있습니다. 이는 많은 머신러닝 애플리케이션에서 원하는 동작이지만, 이는 예상치 못한 결과를 초래할 수 있습니다!

x = random.uniform(random.PRNGKey(0), (1000,), dtype=jnp.float64)
x.dtype

Double 정밀도의 숫자를 사용하려면, 시작 시 jax_enable_x64 구성 변수를 설정해야 합니다.

이를 수행하기 위한 몇가지 방법이 있습니다.

1. JAX_ENABLE_X64=True로 설정하여 64비트 모드를 사용 가능하게 할 수 있습니다.

2. 시작 시에 jax_enable_x64 구성 플래그를 수동으로 설정할 수 있습니다:

# again, this only works on startup!
from jax.config import config
config.update("jax_enable_x64", True)

3. absl.app.run(main)을 사용하여 명령줄 플래그를 파싱할 수 있습니다.

from jax.config import config
config.config_with_absl()

4. absl.app.run(main)를 사용하지 않고 JAX가 absl 파싱을 수행하게 하려면 다음과 사용하면 됩니다:

from jax.config import config
if __name__ == '__main__':
  # calls config.config_with_absl() *and* runs absl parsing
  config.parse_flags_with_absl()

#2-#4는 JAX의 모든 구성 옵션에서 작동합니다.

그런 다음 x64 모드가 활성화되었는지 확인할 수 있습니다.

import jax.numpy as jnp
from jax import random
x = random.uniform(random.PRNGKey(0), (1000,), dtype=jnp.float64)
x.dtype # --> dtype('float64')

주의사항

⚠️ XLA는 모든 백엔드에서 64비트 컨볼루션을 지원하지 않습니다!

🔪 NumPy에서 유래된 여러가지 파생들 (Miscellaneous Divergences from NumPy)

jax.numpy는 Numpy API 동작을 유사하게 동작하도록 설계되었지만, 동작이 다른 코너케이스들이 있습니다. 이러한 많은 경우들은 위 섹션에서 자세히 설명합니다. 여기에는 API가 다르게 동작하는 몇 가지 다른 사례들을 나엻합니다.

• 바이너리 작업의 경우 JAX의 유형 승격 규칙은 NumPy에서 사용하는 규칙과 다소 다릅니다. 자세한 내용은 Type Promotion Semantics를 참조하십시오.

• 안전하지 않은 유형 캐스팅(즉, 대상 dtype이 입력 값을 나타낼 수 없는 캐스팅)를 수행할 때 JAX의 동작은 백엔드에 따라 다를 수 있으며 일반적으로 NumPy의 동작과 다를 수 있습니다. Numpy는 캐스팅 인수를 통해 이러한 시나리오에서 결과를 제어할 수 있습니다(np.ndarray.astype 참조). JAX는 XLA:ConvertElementType의 동작을 직접 상속하는 대신 이러한 구성을 제공하지 않습니다.

여기에 NumPy와 JAX의 안전하지 않은 캐스팅에 따른 다른 결과에 대한 예제입니다.

np.arange(254.0, 258.0).astype('uint8')
array([254, 255,   0,   1], dtype=uint8)

jnp.arange(254.0, 258.0).astype('uint8')
DeviceArray([254, 255, 255, 255], dtype=uint8)

이러한 종류의 불일치는 일반적으로 부동에서 정수 유형으로 또는 그 반대로 극단적인 값을 캐스팅할 때 발생합니다.

Fin.

여기에서 다루지 않은 몇몇 당신을 화나게 하는 원인들을 제보해주시면 해당 튜토리얼 페이지에 반영하겠습니다.