grad의 일: 함수를 받아 그래디언트 함수를 돌려주기

jax.grad는 함수를 받아 새로운 함수를 돌려주는 고차 함수야. 새 함수는 같은 인자를 받지만, 결과는 원래 함수의 그래디언트야.

import jax
import jax.numpy as jnp

def f(x):
    return x ** 2 + 3 * x + 5

# grad(f) 는 새로운 함수
df = jax.grad(f)

print(f(2.0))    # 15.0
print(df(2.0))   # 7.0  (정답: 2x + 3 = 7)
print(df(0.0))   # 3.0
print(df(-1.0))  # 1.0

한 가지 규칙, grad가 받는 함수는 스칼라를 반환해야 해. 여러 출력을 가진 함수의 그래디언트는 수학적으로 Jacobian이며, JAX는 이를 jax.jacrev와 jax.jacfwd로 처리해.

def g(x):
    return jnp.array([x ** 2, x ** 3])

# jax.grad(g)(2.0)  # ❌ TypeError: grad requires scalar output
jacobian = jax.jacrev(g)
print(jacobian(2.0))  # [4., 12.]  ← d/dx [x², x³]

여러 인자에 대한 그래디언트는 기본값으로 첫 번째 인자에 대해서만 구해:

def loss(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y) ** 2)

# 첫 번째 인자 (params) 만
g = jax.grad(loss)(params, x, y)

# 명시적
g = jax.grad(loss, argnums=0)(params, x, y)

# x 에 대한 gradient
g_x = jax.grad(loss, argnums=1)(params, x, y)

# 여러 인자 동시
g_p, g_x = jax.grad(loss, argnums=(0, 1))(params, x, y)

중요한 점, params가 배열일 수도, dict일 수도, 임의 pytree일 수도 있어. grad는 같은 모양의 pytree 그래디언트를 돌려줘:

params = {
    "W1": jnp.zeros((10, 20)),
    "b1": jnp.zeros(20),
    "W2": jnp.zeros((20, 5)),
}

def loss(params, x, y):
    h = jnp.tanh(x @ params["W1"] + params["b1"])
    pred = h @ params["W2"]
    return jnp.mean((pred - y) ** 2)

grads = jax.grad(loss)(params, x, y)
# grads 는 같은 dict 구조: {"W1": ..., "b1": ..., "W2": ...}

🌿 함수형 미분의 우아함

PyTorch에서는 텐서.requires_grad_()와 손실.backward()를 호출하면 텐서.grad가 마법처럼 붙어. JAX에서는 g = jax.grad(loss)(params, x, y)처럼 입력이 들어가고 그래디언트가 나와. 함수형. 어디에도 마법 같은 상태가 없고, 지울 zero_grad()도 없어서 모든 게 보여.

그래디언트의 의미도 짚고 넘어가자. 함수의 출력을 1만큼 늘리려면 입력을 어느 방향으로 얼마나 움직여야 하는가. SGD는 그 반대 방향으로 움직여. 이 직관을 잡아 두면 수학이 덜 헷갈려.

Code

import jax
import jax.numpy as jnp

# A scalar-valued function
def f(x):
    return x ** 2

# grad(f) returns a NEW function that computes df/dx
df = jax.grad(f)

print(f(3.0))    # 9.0
print(df(3.0))   # 6.0 (derivative of x^2 is 2x, evaluated at x=3)

# df is a regular Python function — you can call it, JIT it, etc.
print(df(5.0))   # 10.0

import jax
import jax.numpy as jnp

# Multi-variable function
def loss(params, x, y):
    w, b = params
    pred = jnp.dot(x, w) + b
    return jnp.mean((pred - y) ** 2)

# grad differentiates with respect to the FIRST argument by default
grad_fn = jax.grad(loss)

# Call it: returns gradient with same structure as params
w = jnp.array([1.0, 2.0])
b = jnp.array(0.0)
params = (w, b)
x = jnp.array([[1.0, 0.5], [0.3, 0.8]])
y = jnp.array([1.0, 0.5])

grads = grad_fn(params, x, y)
print(type(grads))       # tuple — same structure as params!
print(grads[0].shape)    # (2,) — gradient of w
print(grads[1].shape)    # () — gradient of b

🌿 함수형 미분의 우아함

Code

External links

Exercise

Progress

댓글 0