🔪 JAX - 尖锐之处 🔪

🔪 JAX - 尖锐的陷阱 🔪#

当你在意大利乡村漫步时，当地人会毫不犹豫地告诉你，JAX 拥有 “纯粹的函数式编程灵魂”。

JAX 是一种用于表达和组合数值程序变换的语言。JAX 还能将数值程序编译为 CPU 或加速器（GPU/TPU）代码。JAX 非常适用于许多数值和科学计算程序，但前提是必须遵循我们下文描述的某些约束。

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

🔪 纯函数#

JAX 的变换和编译设计为仅适用于函数式纯粹的 Python 函数：所有输入数据均通过函数参数传入，所有结果均通过函数返回值输出。如果使用相同的输入调用，纯函数始终会返回相同的结果。

以下是一些非纯函数示例，对于这些函数，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.])))

Executing function
First call:  4.0
Second call:  5.0
Executing function
Third call, different type:  [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.])))

First call:  4.0
Second call:  5.0
Third call, different type:  [14.]

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

First call:  4.0
Saved global:  JitTracer(~float32[])

即便 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.))

50.0

不建议在任何想要 jit 的 JAX 函数或任何控制流原语中使用迭代器。原因是迭代器是一个引入状态以获取下一个元素的 Python 对象。因此，它与 JAX 的函数式编程模型不兼容。在下方的代码中，有一些错误使用迭代器与 JAX 的示例。其中大多数会报错，但有些会产生意外结果。

import jax.numpy as jnp
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

45
0

🔪 原地更新#

在 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)

original array:
[[0. 0. 0.]
 [0. 0. 0.]
 [0. 0. 0.]]
updated array:
[[0. 0. 0.]
 [1. 1. 1.]
 [0. 0. 0.]]

然而，如果我们尝试对 jax.Array 进行原地索引更新，会得到一个错误！(☉_☉)

%xmode Minimal

Exception reporting mode: 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

TypeError: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://jax.net.cn/en/latest/_autosummary/jax.numpy.ndarray.at.html

如果我们尝试进行 __iadd__ 风格的原地更新，会得到与 NumPy 不同的行为！(☉_☉) (☉_☉)

jax_array = jnp.array([10, 20])
jax_array_new = jax_array
jax_array_new += 10
print(jax_array_new)  # `jax_array_new` is rebound to a new value [20, 30], but...
print(jax_array)      # the original value is unmodified as [10, 20] !

numpy_array = np.array([10, 20])
numpy_array_new = numpy_array
numpy_array_new += 10
print(numpy_array_new)  # `numpy_array_new is numpy_array`, and it was updated
print(numpy_array)      # in-place, so both are [20, 30] !

[20 30]
[10 20]
[20 30]
[20 30]

这是因为 NumPy 定义了 __iadd__ 来执行原地修改。相比之下，jax.Array 没有定义 __iadd__，因此 Python 将 jax_array_new += 10 视为 jax_array_new = jax_array_new + 10 的语法糖，重新绑定变量而不修改任何数组。

允许变量的原地修改会使程序分析和变换变得困难。JAX 要求程序是纯函数。

取而代之的是，JAX 提供了使用 .at 属性进行的函数式数组更新。

️⚠️ 在 jit 代码块以及 lax.while_loop 或 lax.fori_loop 中，切片的大小不能是参数值的函数，只能是参数形状的函数——切片的起始索引没有此限制。关于此限制的更多信息，请参阅下方的控制流一节。

数组更新: `x.at[idx].set(y)`#

例如，上述更新可以写为：

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

updated array:
 [[0. 0. 0.]
 [1. 1. 1.]
 [0. 0. 0.]]

与 NumPy 版本不同，JAX 的数组更新函数是“非原地”的。也就是说，更新后的数组作为新数组返回，原始数组不会被修改。

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

original array unchanged:
 [[0. 0. 0.]
 [0. 0. 0.]
 [0. 0. 0.]]

然而，在 jit 编译的代码中，如果 x.at[idx].set(y) 的输入值 x 不再被重用，编译器会优化该数组更新，使其原地执行。

数组与其他操作的更新#

索引数组更新不仅限于覆盖值。例如，我们可以按如下方式执行索引加法：

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)

original array:
[[1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1.]]
new array post-addition:
[[1. 1. 1. 8. 8. 8.]
 [1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 8. 8. 8.]
 [1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 8. 8. 8.]]

有关索引数组更新的更多详细信息，请参阅 .at 属性的文档。

🔪 在类方法中使用 `jax.jit`#

大多数 jax.jit 的示例都涉及修饰独立的 Python 函数，但修饰类中的方法会引入一些复杂性。例如，考虑以下简单的类，其中我们在方法上使用了标准的 jax.jit 注解：

import jax.numpy as jnp
from jax import jit

class CustomClass:
  def __init__(self, x: jnp.ndarray, mul: bool):
    self.x = x
    self.mul = mul

  @jit  # <---- How to do this correctly?
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

然而，当你尝试调用此方法时，这种方法会导致错误：

c = CustomClass(2, True)
c.calc(3)

TypeError: Error interpreting argument to <function CustomClass.calc at 0x783b302c9260> as an abstract array. The problematic value is of type <class '__main__.CustomClass'> and was passed to the function at path self.
This typically means that a jit-wrapped function was called with a non-array argument, and this argument was not marked as static using the static_argnums or static_argnames parameters of jax.jit.

问题在于函数的第一个参数是 self，其类型为 CustomClass，而 JAX 不知道如何处理这种类型。在这种情况下，我们可以使用三种基本策略，我们将在下面讨论它们。

策略 1：JIT 编译的辅助函数#

最直接的方法是创建一个类外部的辅助函数，它可以以常规方式进行 JIT 修饰。例如：

from functools import partial

class CustomClass:
  def __init__(self, x: jnp.ndarray, mul: bool):
    self.x = x
    self.mul = mul

  def calc(self, y):
    return _calc(self.mul, self.x, y)

@partial(jit, static_argnums=0)
def _calc(mul, x, y):
  if mul:
    return x * y
  return y

结果将按预期工作：

c = CustomClass(2, True)
print(c.calc(3))

这种方法的好处是简单、显式，且无需教 JAX 如何处理 CustomClass 类型的对象。然而，你可能希望将所有方法逻辑保持在同一位置。

策略 2：将 `self` 标记为静态#

另一种常见的模式是使用 static_argnums 将 self 参数标记为静态。但这必须谨慎进行，以避免意外结果。你可能会简单地尝试这样做：

class CustomClass:
  def __init__(self, x: jnp.ndarray, mul: bool):
    self.x = x
    self.mul = mul

  # WARNING: this example is broken, as we'll see below. Don't copy & paste!
  @partial(jit, static_argnums=0)
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

如果你调用该方法，它将不再报错：

c = CustomClass(2, True)
print(c.calc(3))

但是有一个陷阱：如果在第一次方法调用后修改了对象，后续的方法调用可能会返回错误的结果：

c.mul = False
print(c.calc(3))  # Should print 3

为什么会这样？当你将一个对象标记为静态时，它实际上会被用作 JIT 内部编译缓存中的字典键，这意味着它的哈希值（即 hash(obj)）、相等性（即 obj1 == obj2）和对象身份（即 obj1 is obj2）都将被假定具有一致的行为。自定义对象的默认 __hash__ 是其对象 ID，因此 JAX 无法知道修改后的对象应该触发重新编译。

你可以通过为对象定义适当的 __hash__ 和 __eq__ 方法来部分解决此问题；例如：

class CustomClass:
  def __init__(self, x: jnp.ndarray, mul: bool):
    self.x = x
    self.mul = mul

  @partial(jit, static_argnums=0)
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

  def __hash__(self):
    return hash((self.x, self.mul))

  def __eq__(self, other):
    return (isinstance(other, CustomClass) and
            (self.x, self.mul) == (other.x, other.mul))

（有关重写 __hash__ 时的要求，请参阅 object.__hash__ 文档了解更多讨论）。

只要你不修改你的对象，这应该能与 JIT 和其他变换正常工作。用作哈希键的对象进行修改会导致几个微妙的问题，这就是为什么例如可变的 Python 容器（如 dict, list）不定义 __hash__，而其不可变对应物（如 tuple）却定义了。

如果你的类依赖于原地修改（例如在其方法内部设置 self.attr = ...），那么你的对象并非真正的“静态”，标记它可能会导致问题。幸运的是，对于这种情况还有另一种选择。

策略 3：将 `CustomClass` 设为 PyTree#

正确 JIT 编译类方法的最灵活方法是将该类型注册为自定义 PyTree 对象；请参阅自定义 PyTree 节点。这让你能精确指定类的哪些组件应视为静态，哪些应视为动态。示例如下：

class CustomClass:
  def __init__(self, x: jnp.ndarray, mul: bool):
    self.x = x
    self.mul = mul

  @jit
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

  def _tree_flatten(self):
    children = (self.x,)  # arrays / dynamic values
    aux_data = {'mul': self.mul}  # static values
    return (children, aux_data)

  @classmethod
  def _tree_unflatten(cls, aux_data, children):
    return cls(*children, **aux_data)

from jax import tree_util
tree_util.register_pytree_node(CustomClass,
                               CustomClass._tree_flatten,
                               CustomClass._tree_unflatten)

这当然涉及更多工作，但它解决了上述简单方法相关的所有问题。

c = CustomClass(2, True)
print(c.calc(3))

c.mul = False  # mutation is detected
print(c.calc(3))

c = CustomClass(jnp.array(2), True)  # non-hashable x is supported
print(c.calc(3))

只要你的 tree_flatten 和 tree_unflatten 函数能正确处理类中的所有相关属性，你无需任何特殊注解，就能直接将此类型的对象作为 JIT 编译函数的参数使用。

🔪 越界索引#

在 NumPy 中，你习惯于在数组越界索引时抛出错误，如下所示：

np.arange(10)[11]

IndexError: index 11 is out of bounds for axis 0 with size 10

然而，在运行于加速器的代码中引发错误可能很困难甚至不可能。因此，JAX 必须为越界索引选择某种非错误的行为（类似于无效浮点算术运算产生 NaN 的方式）。当索引操作是数组索引更新（如 index_add 或类似 scatter 的原语）时，越界索引处的更新将被跳过；当操作是数组索引检索（如 NumPy 索引或类似 gather 的原语）时，索引会被截断（clamp）到数组边界内，因为必须返回某些东西。例如，此索引操作将返回数组的最后一个值：

jnp.arange(10)[11]

Array(9, dtype=int32)

如果你希望更细粒度地控制越界索引的行为，可以使用 ndarray.at 的可选参数；例如：

jnp.arange(10.0).at[11].get()

Array(9., dtype=float32)

jnp.arange(10.0).at[11].get(mode='fill', fill_value=jnp.nan)

Array(nan, dtype=float32)

请注意，由于这种索引检索行为，像 jnp.nanargmin 和 jnp.nanargmax 这样的函数对于由 NaN 组成的切片返回 -1，而 NumPy 则会抛出错误。

另请注意，由于上述两种行为并非互逆，反向模式自动微分（将索引更新变为索引检索，反之亦然）不会保留越界索引的语义。因此，将 JAX 中的越界索引视为未定义行为可能是个好主意。

🔪 非数组输入：NumPy 与 JAX#

NumPy 通常乐于接受 Python 列表或元组作为其 API 函数的输入：

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

np.int64(6)

JAX 则不同，通常会返回有帮助的错误：

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

TypeError: sum requires ndarray or scalar arguments, got <class 'list'> at position 0.

这是一个经过深思熟虑的设计选择，因为将列表或元组传递给跟踪函数可能导致难以察觉的静默性能下降。

例如，考虑以下允许列表输入的 jnp.sum 的宽容版本：

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

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

Array(45, dtype=int32)

输出如我们所料，但这掩盖了底层潜在的性能问题。在 JAX 的跟踪和 JIT 编译模型中，Python 列表或元组中的每个元素都被视为单独的 JAX 变量，并单独处理和推送至设备。这可以在上述 permissive_sum 函数的 jaxpr 中看到：

make_jaxpr(permissive_sum)(x)

{ lambda ; a:i32[] b:i32[] c:i32[] d:i32[] e:i32[] f:i32[] g:i32[] h:i32[] i:i32[]
    j:i32[]. let
    k:i32[] = convert_element_type[new_dtype=int32 weak_type=False] a
    l:i32[1] = broadcast_in_dim k
    m:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b
    n:i32[1] = broadcast_in_dim m
    o:i32[] = convert_element_type[new_dtype=int32 weak_type=False] c
    p:i32[1] = broadcast_in_dim o
    q:i32[] = convert_element_type[new_dtype=int32 weak_type=False] d
    r:i32[1] = broadcast_in_dim q
    s:i32[] = convert_element_type[new_dtype=int32 weak_type=False] e
    t:i32[1] = broadcast_in_dim s
    u:i32[] = convert_element_type[new_dtype=int32 weak_type=False] f
    v:i32[1] = broadcast_in_dim u
    w:i32[] = convert_element_type[new_dtype=int32 weak_type=False] g
    x:i32[1] = broadcast_in_dim w
    y:i32[] = convert_element_type[new_dtype=int32 weak_type=False] h
    z:i32[1] = broadcast_in_dim y
    ba:i32[] = convert_element_type[new_dtype=int32 weak_type=False] i
    bb:i32[1] = broadcast_in_dim ba
    bc:i32[] = convert_element_type[new_dtype=int32 weak_type=False] j
    bd:i32[1] = broadcast_in_dim bc
    be:i32[10] = concatenate[dimension=0] l n p r t v x z bb bd
    bf:i32[] = reduce_sum[axes=(0,) out_sharding=None] be
  in (bf,) }

列表中的每个条目都被作为单独的输入处理，导致跟踪和编译开销随列表大小线性增长。为了防止此类意外，JAX 避免将列表和元组隐式转换为数组。

如果你想将元组或列表传递给 JAX 函数，可以通过先将其显式转换为数组来实现：

jnp.sum(jnp.array(x))

Array(45, dtype=int32)

🔪 随机数#

JAX 的伪随机数生成与 Numpy 有重要区别。有关快速操作指南，请参阅伪随机数。有关更多详细信息，请参阅伪随机数教程。

🔪 控制流#

已移至使用 JIT 的控制流和逻辑运算符。

🔪 动态形状#

在 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))

10.0

如果你尝试将 jax.jit 或其他变换应用于此函数，它会报错：

jax.jit(nansum)(x)

NonConcreteBooleanIndexError: Array boolean indices must be concrete; got bool[5]

See https://jax.net.cn/en/latest/errors.html#jax.errors.NonConcreteBooleanIndexError

问题在于 x_without_nans 的大小取决于 x 中的值，换句话说，其大小是动态的。在 JAX 中，通常可以通过其他方式解决对动态大小数组的需求。例如，在这里可以使用 jnp.where 的三参数形式将 NaN 值替换为零，从而在计算出相同结果的同时避免了动态形状：

@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))

10.0

在出现动态形状数组的其他情况下，也可以使用类似的技巧。

🔪 调试 NaN 和 Inf#

使用 jax_debug_nans 和 jax_debug_infs 标志来查找函数和梯度中 NaN/Inf 值的来源。请参阅 JAX 调试标志。

🔪 双精度 (64位)#

目前，JAX 默认强制使用单精度数字，以减轻 Numpy API 将操作数激进提升为 double 的倾向。对于许多机器学习应用，这是预期的行为，但它可能会让你措手不及！

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

/tmp/ipykernel_2210/1258726447.py:1: UserWarning: Explicitly requested dtype float64 is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  x = random.uniform(random.key(0), (1000,), dtype=jnp.float64)

dtype('float32')

要使用双精度数字，你需要在启动时设置 jax_enable_x64 配置变量。

有几种方法可以做到这一点：

你可以通过设置环境变量 JAX_ENABLE_X64=True 来启用 64 位模式。

你可以在启动时手动设置 jax_enable_x64 配置标志：

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

你可以使用 absl.app.run(main) 解析命令行标志：
```
import jax
jax.config.config_with_absl()
```

如果你希望 JAX 为你运行 absl 解析，即你不想手动使用 absl.app.run(main)，你可以使用：

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

请注意，#2-#4 适用于 JAX 的任何配置选项。

然后我们可以确认 x64 模式已启用，例如：

import jax
import jax.numpy as jnp
from jax import random

jax.config.update("jax_enable_x64", True)
x = random.uniform(random.key(0), (1000,), dtype=jnp.float64)
x.dtype # --> dtype('float64')

注意事项#

⚠️ XLA 不支持所有后端上的 64 位卷积！

🔪 与 NumPy 的其他差异#

尽管 jax.numpy 尽一切努力复制 NumPy API 的行为，但仍存在行为不同的边缘情况。许多此类情况已在上述章节中详细讨论；在此我们列出 API 差异的其他几个已知点。

对于二元运算，JAX 的类型提升规则与 NumPy 使用的规则略有不同。详见类型提升语义。
执行不安全的类型转换时（即目标 dtype 无法表示输入值的转换），JAX 的行为可能取决于后端，并且通常可能与 NumPy 的行为不同。NumPy 允许通过 casting 参数控制这些场景中的结果（参见 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')
Array([254, 255, 255, 255], dtype=uint8)
```
这种不匹配通常出现在将极端数值从浮点类型转换到整数类型或反之时。

在对次正规（subnormal）浮点数进行运算时，JAX 在某些后端上使用“清零”（flush-to-zero）语义。例如：

>>> import jax.numpy as jnp
>>> subnormal = jnp.float32(1E-45)
>>> subnormal  # subnormals are representable
Array(1.e-45, dtype=float32)
>>> subnormal + 0  # but are flushed to zero within operations
Array(0., dtype=float32)

次正规值的详细操作语义通常会根据后端而异。

🔪 教程中涵盖的尖锐陷阱#

使用 JIT 的控制流和逻辑运算符讨论了如何处理 jit 对 Python 控制流和逻辑运算符使用施加的约束。
有状态计算提供了一些关于如何正确处理 JAX 程序中状态的建议，鉴于 JAX 变换只能应用于纯函数。

结束。#

如果这里没有涵盖的内容让你感到困扰或沮丧，请告诉我们，我们将扩充这些入门建议！