jax.ensure_compile_time_eval

jax.ensure_compile_time_eval#

jax.ensure_compile_time_eval()[source]#

上下文管理器，确保在跟踪/编译时进行评估（或错误）。

一些 JAX API（如 jax.jit() 和 jax.lax.scan()）涉及分阶段，即延迟数值表达式的评估（如 jax.numpy 函数应用），以便在评估相应的 Python 表达式时不急切地执行这些计算，而是在其他情况下执行它们，例如在经过优化的编译之后。但是这种延迟可能是不可取的。例如，可能需要数值来评估 Python 控制流，因此它们的评估不能被延迟。作为另一个例子，为了性能原因，确保编译时评估（或“常量折叠”）可能是有益的。

此上下文管理器确保 JAX 计算被急切地评估。如果无法进行急切评估，则会引发 ConcretizationTypeError。

这是一个人为的例子

import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  with jax.ensure_compile_time_eval():
    y = jnp.sin(3.0)
    z = jnp.sin(y)
    z_positive = z > 0
  if z_positive:  # z_positive is usable in Python control flow
    return jnp.sin(x)
  else:
    return jnp.cos(x)

这是一个来自 google/jax#3974 的真实世界示例

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

@jax.jit
def jax_fn(x):
  with jax.ensure_compile_time_eval():
    y = random.randint(random.key(0), (1000,1000), 0, 100)
  y2 = y @ y
  x2 = jnp.sum(y2) * x
  return x2

类似的行为通常可以通过将常量表达式“提升”到相应的分阶段 API 之外来实现

y = random.randint(random.key(0), (1000,1000), 0, 100)

@jax.jit
def jax_fn(x):
  y2 = y @ y
  x2 = jnp.sum(y2)*x
  return x2

但在某些情况下，使用此上下文管理器可能更方便。

内容