TF2JAX

CI 状态 pypi

TF2JAX 是一个实验性库，用于将 [TensorFlow] 函数/图转换为 [JAX] 函数。

具体来说，它旨在将 tf.function 转换，例如：

@tf.function
def tf_fn(x):
  return tf.sin(tf.cos(x))

转换为等效的以下 JAX 代码的 Python 函数：

def jax_fn(x):
  return jnp.sin(jnp.cos(x))

用户可以对转换后的函数应用额外的 JAX 变换（如 jit、grad、vmap、make_jaxpr 等），就像对其他用 JAX 编写的代码一样。

[目录]

安装

您可以通过 PyPI 安装 TF2JAX 的最新发布版本：

pip install tf2jax

或者您可以从 GitHub 安装最新的开发版本：

pip install git+https://github.com/google-deepmind/tf2jax.git

动机

TF2JAX 使现有的 TensorFlow 函数和模型（包括 SavedModel 和 TensorFlow Hub）能够在 JAX 代码库中重用和/或微调。转换过程对用户来说是透明的，这对调试和内省很有用。

这也为 JAX 用户提供了一种途径，将通过 jax2tf.convert 序列化的 JAX 函数重新集成到他们现有的 JAX 代码库中。

请参阅末尾的章节，了解与 jax2tf.call_tf 提供的替代方法的比较。

免责声明

这是一个实验性代码，API 可能不稳定，目前不对使用它提供任何保证。我们强烈建议您彻底测试生成的 JAX 函数，以确保它们满足您的要求。

快速开始

本文档的其余部分假设有以下导入：

import jax
import jax.numpy as jnp
import numpy as np
import tensorflow as tf  # 假设这是 v2。
import tf2jax

使用 convert API 和 Sonnet v2 MLP 的示例。

import sonnet.v2 as snt

model = snt.nets.MLP((64, 10,))

@tf.function
def forward(x):
  return model(x)

x = np.random.normal(size=(128, 16)).astype(np.float32)

# TF -> JAX，jax_params 是 MLP 的网络参数
jax_func, jax_params = tf2jax.convert(forward, np.zeros_like(x))

# 调用 JAX，同时返回更新后的 jax_params（如变量、批归一化统计信息）
jax_outputs, jax_params = jax_func(jax_params, x)

tf2jax.convert 的签名为 convert(fn: tf.Function, *args, **kwargs)，其中 fn(*args, **kwargs) 用于追踪函数 fn 并生成相应的 tf.GraphDef。这里的 zeros_like 不是必需的，只是用来演示 JAX 函数并没有记忆输出。

纯函数示例

如果您的函数是纯函数，即它不捕获任何变量，那么您可以使用 tf2jax.convert_functional 从转换后函数的输入和输出中删除参数。

@tf.function
def forward(x):
  return tf.sin(tf.cos(x))

jax_func = tf2jax.convert_functional(forward, np.zeros_like(x))
jax_outputs = jax_func(x)

随机性和 PRNG 键

使用随机操作的 TensorFlow 函数将被转换为一个 JAX 函数，该函数将 PRNG 键作为仅关键字参数。如果需要 PRNG 键但未提供，TF2JAX 将大声抱怨。

jax_outputs, jax_params = jax_func(jax_params, x, rng=jax.random.PRNGKey(42))

自定义梯度

自定义梯度支持仍处于高度实验阶段，请报告任何错误。

@tf.function
@tf.custom_gradient
def forward(x):
  e = tf.exp(x)
  def grad(dy):
    return dy * tf.sin(x) + e  # 这是故意错误的梯度。
  return tf.reduce_sum(e), grad

with tf2jax.override_config("convert_custom_gradient", True):
  jax_func = tf2jax.convert_functional(forward, np.zeros_like(x))

jax_grads = jax.grad(jax_func)(x)

序列化格式支持

SavedModel

SavedModel 是序列化 TF2 函数的首选格式。

model = tf.Module()
model.f = forward
model.f(x)  # 虚拟调用。
tf.saved_model.save(model, "/tmp/blah")

restored = tf.saved_model.load("/tmp/blah")
jax_func, jax_params = tf2jax.convert(restored.f, np.zeros_like(x))

如果恢复的函数具有明确的签名，即在导出之前只被追踪过一次。那么 TF2JAX 可以直接从其 GraphDef 转换函数，而无需再次追踪。

jax_func, jax_params = tf2jax.convert_from_restored(restored.f)

TF-Hub

支持（传统的 TF1）TF-Hub 格式，只需少量样板代码。

import tensorflow_hub as hub

hub_model = hub.load("/tmp/blah")
jax_func, jax_params = tf2jax.convert(tf.function(hub_model), tf.zeros_like(x))
jax_outputs, updated_jax_params = jax_func(jax_params, x)

JAX 到 TensorFlow 再回到 JAX

tf2jax.convert_functional 可以将 jax2tf.convert 的输出转换回 JAX 代码。

# 一些 JAX 函数。
def forward(*inputs):
  ...

# JAX -> TF
tf_func = jax2tf.convert(forward)

# JAX -> TF -> JAX
jax_func = tf2jax.convert_functional(tf.function(tf_func), *tree.map_structure(np.zeros_like, inputs))

# JAX -> TF -> SavedModel -> TF
model = tf.Module()
model.f = tf.function(tf_func)
model.f(*tree.map_structure(tf.zeros_like, inputs))  # 虚拟调用。
tf.saved_model.save(model, "/tmp/blah")
restored = tf.saved_model.load("/tmp/blah")

# JAX -> TF -> SavedModel -> TF -> JAX
jax_too_func = tf2jax.convert_functional(restored.f, *tree.map_structure(np.zeros_like, inputs))

附加配置

TF2JAX 的行为可以通过 tf2jax.update_config 全局配置，或通过上下文管理器 tf2jax.override_config 在局部配置。

严格的形状和数据类型检查

默认情况下，TF2JAX 将断言转换后函数的输入形状与原始函数的输入形状兼容。这是因为一些函数具有依赖于形状的行为，在转换后会静默地返回不正确的输出，例如某些批归一化实现。

jax_func = tf2jax.convert_functional(forward, np.zeros((10, 5), np.float32))

# 这将引发错误。
jax_func(np.zeros((20, 5), np.float32))

# 这不会。
with tf2jax.override_config("strict_shape_check", False):
  jax_func(np.zeros((20, 5), np.float32))

目前不检查输入数据类型，但这可能在将来发生变化。

jax_func = tf2jax.convert_functional(forward, np.zeros((10, 5), np.float32))

# 这不会引发错误。
jax_func(np.zeros((20, 5), np.int32))

# 这会。
with tf2jax.override_config("strict_dtype_check", True):
  jax_func(np.zeros((20, 5), np.int32))

将常量转换为 bfloat16

TF2JAX 允许用户使用与原始函数不同精度的参数和输入来追踪转换后的函数，例如使用 bfloat16 而不是 float32。为了帮助实现这一点，配置 force_const_float32_to_bfloat16 和 force_const_float64_to_bfloat16 可以用来强制将原始函数中的浮点常量转换为 bfloat16 精度，以避免意外的类型提升。

@tf.function
def forward(x):
  return x + tf.constant(3.14, dtype=tf.float32)

with tf2jax.override_config("force_const_float32_to_bfloat16", True):
  jax_func = tf2jax.convert_functional(forward, np.zeros_like(x))
jax_bf16_outputs = jax_func(jnp.asarray(x, jnp.bfloat16))

禁用 PreventGradient

如果使用 jax2tf.convert(..., with_gradient=False) 生成初始 TF 函数（并可能导出为 SavedModel），那么 TF2JAX 将尊重插入的 tf.raw_ops.PreventGradient 操作，并在计算梯度时引发 LookupError。

通过将配置 raise_on_prevent_gradient 设置为 false（默认为 true），可以禁用此行为，这样 TF2JAX 将只记录警告，但允许计算梯度，就好像 PreventGradient 操作不存在一样。

@tf.function
def prevent(x):
  return tf.raw_ops.PreventGradient(input=x * x, message="prevented")

jax_func = tf2jax.convert_functional(prevent, 0.0)
jax.grad(jax_func)(3.14)  # 引发 LookupError。

with tf2jax.config.override_config("raise_on_prevent_gradient", False):
  jax_func = tf2jax.convert_functional(prevent, 0.0)
g = jax.grad(jax_func)(3.14)  # 返回 6.28。

推断累积归约

如果 infer_cumulative_reduction_from_jax2tf 标志为 true（默认），则 TF2JAX 将尝试从 JAX2TF 生成的 reduce_window 操作中推断 cummax、cummin、cumprod 和 cumsum 操作。这提供了更好的性能，因为这些操作的 reduce_window 实现在 CPU 和 GPU 后端上具有 O(N^2) 的复杂度，并且由于积极的常量折叠可能会导致长时间的编译时间。

有关更多背景信息，请参阅 [jax2tf_cumulative_reduction]。

JAX2TF 原生序列化和 XlaCallModule

从 JAX v0.4.7 开始，jax2tf.convert 的首选操作模式（即将成为默认模式）是原生序列化，其中目标函数被降级为 [StableHLO] 并包装在单个 TensorFlow 操作 XlaCallModule 中。新的原生序列化格式将更忠实地再现目标函数的语义，代价是下游处理的灵活性有所降低，因为计算图不再以tf.Graph的形式暴露。

从v0.3.4版本开始，TF2JAX支持XlaCallModule。

然而，由于这使用了一个自定义JAX原语，旨在封装XlaCallModule中的[StableHLO]载荷，它不具备JAX规则来进行诸如批处理和微分等转换（但不限于这些）。

微分：通过在序列化时使用jax2tf.convert(..., with_gradient=True)请求自定义梯度，仍然支持序列化函数的一阶导数。这是默认行为。
批处理：jax.vmap将会失败，不过用户可能可以通过jax.lax.map简单地复制所需行为，尽管性能会较差。

平台特异性

原生序列化的JAX程序是平台特定的。在非目标平台上执行原生序列化程序会引发ValueError，例如：

ValueError: 不支持的后端：`cpu`不在`('tpu',)`中。

这与XlaCallModule的行为一致。

用户可以通过配置标志禁用此检查，但得到的程序可能会变慢或完全无法执行。

with tf2jax.override_config("xlacallmodule_strict_checks", False):
  jax_func(np.zeros((20, 5), np.float32))

限制

目前，仅支持部分TensorFlow算子，某些算子并非所有功能都受支持。代码会快速失败。根据需求添加对其他TensorFlow算子的支持。请通过Github问题提交您的请求或发送拉取请求。

在某些情况下，生成的JAX代码可能在性能和数值输出方面与TensorFlow代码不等效。我们的目标是最小化模型端点在后者方面的差异，然后再改进性能。

TF2控制流受支持，但有一些限制，例如对于while循环，cond和body函数不能有副作用，如给变量赋值。

不支持TF1控制流。

替代方案

`jax2tf.call_tf`

jax2tf现在还提供了实验性的call_tf函数，允许JAX调用TensorFlow函数。对于编译代码，这通过将TensorFlow函数阶段化到XLA来实现。

根据2022-07-22的[jax2tf文档]：

call_tf函数允许JAX函数调用TensorFlow函数。这些函数可以在JAX计算的任何位置调用，包括在jax.jit、jax.pmap、jax.xmap等暂存上下文中，或在JAX的控制流原语内部。在非暂存上下文中，TensorFlow函数以急切模式调用。目前，这些函数仅支持反向模式自动微分（不支持前向模式自动微分，也不支持vmap）。

call_tf的优势在于它隐式覆盖了所有TensorFlow算子，并通过在急切执行期间推迟到TensorFlow，在编译代码时推迟到XLA来支持custom_gradient。

缺点是它只支持有限的JAX转换（jit、grad、pmap、remat），否则对JAX来说就像一个"黑盒"（例如，不支持vmap，也不支持自定义转换）。如果要在call_tf之后进行jit编译，TensorFlow函数必须可编译为XLA。