在 JAX 中编写自定义 Jaxpr 解释器#
JAX 提供了几个可组合的函数转换(jit
、grad
、vmap
等),这些转换使编写简洁、加速的代码成为可能。
在这里,我们展示了如何通过编写自定义 Jaxpr 解释器来向系统添加自己的函数转换。并且我们将免费获得与所有其他转换的可组合性。
此示例使用 JAX 内部 API,这些 API 可能会随时更改。API 文档中未列出的任何内容都应视为内部内容。
import jax
import jax.numpy as jnp
from jax import jit, grad, vmap
from jax import random
JAX 在做什么?#
JAX 提供了一个类似 NumPy 的数值计算 API,可以按原样使用,但 JAX 的真正强大之处在于可组合的函数转换。以 jit
函数转换为例,它接收一个函数并返回一个语义上相同的函数,但由 XLA 为加速器延迟编译。
x = random.normal(random.key(0), (5000, 5000))
def f(w, b, x):
return jnp.tanh(jnp.dot(x, w) + b)
fast_f = jit(f)
当我们调用 fast_f
时,会发生什么?JAX 会跟踪函数并构建一个 XLA 计算图。然后,该图会被 JIT 编译并执行。其他转换的工作方式与此类似,它们首先跟踪函数并以某种方式处理输出跟踪。要了解有关 Jax 跟踪机制的更多信息,您可以参考 README 中的 “工作原理” 部分。
Jaxpr 跟踪器#
Jax 中一个特别重要的追踪器是 Jaxpr 追踪器,它将操作记录到 Jaxpr(Jax 表达式)中。Jaxpr 是一种数据结构,可以像迷你函数式编程语言一样进行评估,因此 Jaxpr 是函数转换的有用中间表示。
为了初步了解 Jaxpr,可以考虑 make_jaxpr
转换。 make_jaxpr
本质上是一个“漂亮打印”转换:它将函数转换为一个函数,该函数在给定示例参数的情况下,会生成其计算的 Jaxpr 表示。 make_jaxpr
对于调试和内省很有用。让我们用它来查看一些示例 Jaxpr 的结构。
def examine_jaxpr(closed_jaxpr):
jaxpr = closed_jaxpr.jaxpr
print("invars:", jaxpr.invars)
print("outvars:", jaxpr.outvars)
print("constvars:", jaxpr.constvars)
for eqn in jaxpr.eqns:
print("equation:", eqn.invars, eqn.primitive, eqn.outvars, eqn.params)
print()
print("jaxpr:", jaxpr)
def foo(x):
return x + 1
print("foo")
print("=====")
examine_jaxpr(jax.make_jaxpr(foo)(5))
print()
def bar(w, b, x):
return jnp.dot(w, x) + b + jnp.ones(5), x
print("bar")
print("=====")
examine_jaxpr(jax.make_jaxpr(bar)(jnp.ones((5, 10)), jnp.ones(5), jnp.ones(10)))
foo
=====
invars: [Var(id=140091243167616):int32[]]
outvars: [Var(id=140091243165504):int32[]]
constvars: []
equation: [Var(id=140091243167616):int32[], 1] add [Var(id=140091243165504):int32[]] {}
jaxpr: { lambda ; a:i32[]. let b:i32[] = add a 1 in (b,) }
bar
=====
invars: [Var(id=140091243526848):float32[5,10], Var(id=140091243526912):float32[5], Var(id=140091243526976):float32[10]]
outvars: [Var(id=140091243527232):float32[5], Var(id=140091243526976):float32[10]]
constvars: []
equation: [Var(id=140091243526848):float32[5,10], Var(id=140091243526976):float32[10]] dot_general [Var(id=140091243527040):float32[5]] {'dimension_numbers': (((1,), (0,)), ((), ())), 'precision': None, 'preferred_element_type': dtype('float32')}
equation: [Var(id=140091243527040):float32[5], Var(id=140091243526912):float32[5]] add [Var(id=140091243527104):float32[5]] {}
equation: [1.0] broadcast_in_dim [Var(id=140091243527168):float32[5]] {'shape': (5,), 'broadcast_dimensions': ()}
equation: [Var(id=140091243527104):float32[5], Var(id=140091243527168):float32[5]] add [Var(id=140091243527232):float32[5]] {}
jaxpr: { lambda ; a:f32[5,10] b:f32[5] c:f32[10]. let
d:f32[5] = dot_general[
dimension_numbers=(([1], [0]), ([], []))
preferred_element_type=float32
] a c
e:f32[5] = add d b
f:f32[5] = broadcast_in_dim[broadcast_dimensions=() shape=(5,)] 1.0
g:f32[5] = add e f
in (g, c) }
jaxpr.invars
- Jaxpr 的invars
是 Jaxpr 输入变量的列表,类似于 Python 函数中的参数。jaxpr.outvars
- Jaxpr 的outvars
是 Jaxpr 返回的变量。每个 Jaxpr 都有多个输出。jaxpr.constvars
-constvars
是一个变量列表,它们也是 Jaxpr 的输入,但对应于跟踪中的常量(稍后我们将详细介绍这些常量)。jaxpr.eqns
- 方程式的列表,本质上是 let 绑定。每个方程式都是一个输入变量列表、一个输出变量列表和一个原始操作,用于评估输入以产生输出。每个方程式还有一个params
,即参数的字典。
总而言之,Jaxpr 封装了一个简单的程序,可以用输入进行评估以产生输出。我们稍后将详细介绍如何执行此操作。现在需要注意的是,Jaxpr 是一种可以以我们想要的方式进行操作和评估的数据结构。
为什么 Jaxpr 有用?#
Jaxpr 是简单的程序表示,易于转换。并且由于 Jax 允许我们从 Python 函数中分阶段提取 Jaxpr,因此它为我们提供了一种转换用 Python 编写的数值程序的方法。
你的第一个解释器:invert
#
让我们尝试实现一个简单的函数“反转器”,它接收原始函数的输出并返回产生这些输出的输入。现在,让我们专注于由其他可逆一元函数组成的简单的一元函数。
目标
def f(x):
return jnp.exp(jnp.tanh(x))
f_inv = inverse(f)
assert jnp.allclose(f_inv(f(1.0)), 1.0)
我们将通过以下方式实现这一点:(1)将 f
跟踪到 Jaxpr 中,然后(2)反向解释 Jaxpr。在反向解释 Jaxpr 时,对于每个方程式,我们将在表中查找原始操作的反函数并应用它。
1. 跟踪函数#
让我们使用 make_jaxpr
将函数跟踪到 Jaxpr 中。
# Importing Jax functions useful for tracing/interpreting.
from functools import wraps
from jax import core
from jax import lax
from jax._src.util import safe_map
jax.make_jaxpr
返回一个封闭的 Jaxpr,它是一个与跟踪中的常量(literals
)捆绑在一起的 Jaxpr。
def f(x):
return jnp.exp(jnp.tanh(x))
closed_jaxpr = jax.make_jaxpr(f)(jnp.ones(5))
print(closed_jaxpr.jaxpr)
print(closed_jaxpr.literals)
{ lambda ; a:f32[5]. let b:f32[5] = tanh a; c:f32[5] = exp b in (c,) }
[]
2. 评估 Jaxpr#
在我们编写自定义 Jaxpr 解释器之前,让我们首先实现“默认”解释器 eval_jaxpr
,它按原样评估 Jaxpr,计算与原始未转换的 Python 函数相同的数值。
为此,我们首先创建一个环境来存储每个变量的值,并在评估 Jaxpr 中的每个方程式时更新环境。
def eval_jaxpr(jaxpr, consts, *args):
# Mapping from variable -> value
env = {}
def read(var):
# Literals are values baked into the Jaxpr
if type(var) is core.Literal:
return var.val
return env[var]
def write(var, val):
env[var] = val
# Bind args and consts to environment
safe_map(write, jaxpr.invars, args)
safe_map(write, jaxpr.constvars, consts)
# Loop through equations and evaluate primitives using `bind`
for eqn in jaxpr.eqns:
# Read inputs to equation from environment
invals = safe_map(read, eqn.invars)
# `bind` is how a primitive is called
outvals = eqn.primitive.bind(*invals, **eqn.params)
# Primitives may return multiple outputs or not
if not eqn.primitive.multiple_results:
outvals = [outvals]
# Write the results of the primitive into the environment
safe_map(write, eqn.outvars, outvals)
# Read the final result of the Jaxpr from the environment
return safe_map(read, jaxpr.outvars)
closed_jaxpr = jax.make_jaxpr(f)(jnp.ones(5))
eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, jnp.ones(5))
[Array([2.1416876, 2.1416876, 2.1416876, 2.1416876, 2.1416876], dtype=float32)]
请注意,即使原始函数没有返回扁平列表,eval_jaxpr
也始终会返回一个扁平列表。
此外,此解释器不处理高阶原始操作(如 jit
和 pmap
),在本指南中我们不会介绍这些操作。您可以参考 core.eval_jaxpr
(链接)以查看此解释器未涵盖的边缘情况。
自定义 inverse
Jaxpr 解释器#
一个 inverse
解释器与 eval_jaxpr
看起来并没有太大区别。我们首先设置注册表,它将原始操作映射到其反函数。然后,我们将编写一个自定义解释器,在注册表中查找原始操作。
事实证明,此解释器也将类似于反向模式自动微分中使用的“转置”解释器 此处可见。
inverse_registry = {}
现在,我们将为一些原始操作注册反函数。按照惯例,Jax 中的原始操作以 _p
结尾,许多常用的原始操作位于 lax
中。
inverse_registry[lax.exp_p] = jnp.log
inverse_registry[lax.tanh_p] = jnp.arctanh
inverse
将首先跟踪函数,然后自定义解释 Jaxpr。让我们设置一个简单的框架。
def inverse(fun):
@wraps(fun)
def wrapped(*args, **kwargs):
# Since we assume unary functions, we won't worry about flattening and
# unflattening arguments.
closed_jaxpr = jax.make_jaxpr(fun)(*args, **kwargs)
out = inverse_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, *args)
return out[0]
return wrapped
现在,我们只需要定义 inverse_jaxpr
,它将向后遍历 Jaxpr 并尽可能反转原始操作。
def inverse_jaxpr(jaxpr, consts, *args):
env = {}
def read(var):
if type(var) is core.Literal:
return var.val
return env[var]
def write(var, val):
env[var] = val
# Args now correspond to Jaxpr outvars
safe_map(write, jaxpr.outvars, args)
safe_map(write, jaxpr.constvars, consts)
# Looping backward
for eqn in jaxpr.eqns[::-1]:
# outvars are now invars
invals = safe_map(read, eqn.outvars)
if eqn.primitive not in inverse_registry:
raise NotImplementedError(
f"{eqn.primitive} does not have registered inverse.")
# Assuming a unary function
outval = inverse_registry[eqn.primitive](*invals)
safe_map(write, eqn.invars, [outval])
return safe_map(read, jaxpr.invars)
就是这样!
def f(x):
return jnp.exp(jnp.tanh(x))
f_inv = inverse(f)
assert jnp.allclose(f_inv(f(1.0)), 1.0)
重要的是,您可以跟踪 Jaxpr 解释器。
jax.make_jaxpr(inverse(f))(f(1.))
{ lambda ; a:f32[]. let b:f32[] = log a; c:f32[] = atanh b in (c,) }
这就是向系统添加新转换所需的一切,并且您可以免费获得与所有其他转换的组合!例如,我们可以将 jit
、vmap
和 grad
与 inverse
一起使用!
jit(vmap(grad(inverse(f))))((jnp.arange(5) + 1.) / 5.)
Array([-3.1440797, 15.584931 , 2.2551253, 1.3155028, 1. ], dtype=float32, weak_type=True)
读者练习#
处理具有多个参数的原始操作,其中输入部分已知,例如
lax.add_p
、lax.mul_p
。处理
xla_call
和xla_pmap
原始操作,它们无法与按原样编写的eval_jaxpr
和inverse_jaxpr
一起使用。