JAX 内部：原语#
JAX 原语简介#
JAX 原语是 JAX 程序的计算基本单元。本文档解释了 JAX 原语必须支持的接口，以便 JAX 能够执行其所有变换（这不是一个操作指南）。
例如，乘加运算可以实现为基于底层 jax.lax.* 原语（它们类似于 XLA 运算符包装器），或者 jax.extend.core.Primitive("multiply_add")，如下面的示例所示。
JAX 能够接收此类原始操作的序列，并通过其可组合的 Python 函数变换来转换它们，例如 jax.jit()、jax.grad() 和 jax.vmap()。JAX 以一种 *JAX 可追踪* 的方式实现这些变换。这意味着当执行 Python 函数时，它仅对数据应用以下操作：
数据属性检查： 数据信息，例如形状或类型；或者
JAX 原语： 这些是本教程中涵盖的 JAX 特殊操作。
JAX 原语知道如何操作具体数据值和抽象 JAX 值。*JAX 可追踪函数* 可以用抽象参数由 JAX 调用。例如，JAX 抽象值 — ShapedArray(float32[2,2]) — 捕获值的类型和形状，但不是具体的数据值。
JAX 变换后的函数本身必须是 JAX 可追踪函数，*以确保这些变换是可组合的*，例如 jax.jit(jax.jacfwd(jax.grad(f)))。
JAX 提供了对应大多数 XLA 操作的预定义原语，包括加法、矩阵乘法、正弦、余弦和索引。
此外，JAX 还提供了基于 JAX 原语实现的 NumPy 函数。这意味着*使用 JAX 实现的 NumPy 的 Python 程序是 JAX 可追踪的，因此是可变换的*。通过用 JAX 原语实现，其他库也可以变成 JAX 可追踪的。
此外，JAX 原语集合是可扩展的，因此您不必用预定义的 JAX 原语重新实现函数，而是可以定义一个封装函数行为的新原语。
考虑以下示例：您想为 JAX 添加一个三参数的乘加函数支持，该函数在数学上定义为 multiply_add(x, y, z) = x * y + z。此函数逐点操作 3 个相同形状的浮点值张量。您可以通过以下方式实现：
使用现有的 JAX 原语#
定义新函数的最简单方法是用 JAX 原语编写它们，或者用其他已经用 JAX 原语编写的函数编写它们，例如 jax.lax() 模块中定义的那些。
from jax._src.lax import lax
from jax._src import api
def multiply_add_lax(x, y, z):
"""Implementation of multiply-add using the `jax.lax` primitives."""
return lax.add(lax.mul(x, y), z)
def square_add_lax(a, b):
"""A square-add function using the newly defined multiply-add."""
return multiply_add_lax(a, a, b)
print("square_add_lax = ", square_add_lax(2., 10.))
# Differentiate w.r.t. the first argument
print("grad(square_add_lax) = ", api.grad(square_add_lax, argnums=0)(2.0, 10.))
square_add_lax = 14.0
grad(square_add_lax) = 4.0
要理解 JAX 在内部如何使用原语，请添加一些用于追踪函数调用的辅助函数。
#@title Helper functions (execute this cell)
import functools
import traceback
_indentation = 0
def _trace(msg=None):
"""Print a message at current indentation."""
if msg is not None:
print(" " * _indentation + msg)
def _trace_indent(msg=None):
"""Print a message and then indent the rest."""
global _indentation
_trace(msg)
_indentation = 1 + _indentation
def _trace_unindent(msg=None):
"""Unindent then print a message."""
global _indentation
_indentation = _indentation - 1
_trace(msg)
def trace(name):
"""A decorator for functions to trace arguments and results."""
def trace_func(func): # pylint: disable=missing-docstring
def pp(v):
"""Print certain values more succinctly"""
vtype = str(type(v))
if "jax._src.xla_bridge._JaxComputationBuilder" in vtype:
return "<JaxComputationBuilder>"
elif "jaxlib._jax_.XlaOp" in vtype:
return "<XlaOp at 0x{:x}>".format(id(v))
elif ("partial_eval.JaxprTracer" in vtype or
"batching.BatchTracer" in vtype or
"ad.JVPTracer" in vtype):
return "Traced<{}>".format(v.aval)
elif isinstance(v, tuple):
return "({})".format(pp_values(v))
else:
return str(v)
def pp_values(args):
return ", ".join([pp(arg) for arg in args])
@functools.wraps(func)
def func_wrapper(*args):
_trace_indent("call {}({})".format(name, pp_values(args)))
res = func(*args)
_trace_unindent("|<- {} = {}".format(name, pp(res)))
return res
return func_wrapper
return trace_func
class expectNotImplementedError(object):
"""Context manager to check for NotImplementedError."""
def __enter__(self): pass
def __exit__(self, type, value, tb):
global _indentation
_indentation = 0
if type is NotImplementedError:
print("\nFound expected exception:")
traceback.print_exc(limit=3)
return True
elif type is None: # No exception
assert False, "Expected NotImplementedError"
else:
return False
您可以使用 jax.numpy 中的函数，而不是直接使用 jax.lax() 原语。
import jax.numpy as jnp
import numpy as np
@trace("multiply_add_numpy")
def multiply_add_numpy(x, y, z):
return jnp.add(jnp.multiply(x, y), z)
@trace("square_add_numpy")
def square_add_numpy(a, b):
return multiply_add_numpy(a, a, b)
print("\nNormal evaluation:")
print("square_add_numpy = ", square_add_numpy(2., 10.))
print("\nGradient evaluation:")
print("grad(square_add_numpy) = ", api.grad(square_add_numpy)(2.0, 10.))
Normal evaluation:
call square_add_numpy(2.0, 10.0)
call multiply_add_numpy(2.0, 2.0, 10.0)
|<- multiply_add_numpy = 14.0
|<- square_add_numpy = 14.0
square_add_numpy = 14.0
Gradient evaluation:
call square_add_numpy(LinearizeTracer<~float32[]>, 10.0)
call multiply_add_numpy(LinearizeTracer<~float32[]>, LinearizeTracer<~float32[]>, 10.0)
|<- multiply_add_numpy = LinearizeTracer<~float32[]>
|<- square_add_numpy = LinearizeTracer<~float32[]>
grad(square_add_numpy) = 4.0
请注意，在计算 jax.grad() 的过程中，JAX 使用特殊参数 ConcreteArray(...)（本 Colab 文档稍后将详细介绍）调用 square_add_numpy 和 multiply_add_numpy。务必记住，JAX 可追踪函数不仅必须能够处理具体参数，还必须能够处理 JAX 可能用于抽象函数执行的特殊抽象参数。
只要函数是用 JAX 原语编写的，JAX 可追踪性属性就能得到满足。
定义新的 JAX 原语#
添加乘加函数支持的正确方法是使用现有的 JAX 原语，如上所示。但是，为了演示 JAX 原语的工作原理，请假装您想为乘加功能向 JAX 添加一个新原语。
from jax.extend import core
multiply_add_p = core.Primitive("multiply_add") # Create the primitive
@trace("multiply_add_prim")
def multiply_add_prim(x, y, z):
"""The JAX-traceable way to use the JAX primitive.
Note that the traced arguments must be passed as positional arguments
to `bind`.
"""
return multiply_add_p.bind(x, y, z)
@trace("square_add_prim")
def square_add_prim(a, b):
"""A square-add function implemented using the new JAX-primitive."""
return multiply_add_prim(a, a, b)
如果您尝试调用新定义的函数，将会收到一个错误，因为您还没有告诉 JAX 新原语的语义。
with expectNotImplementedError():
square_add_prim(2., 10.)
call square_add_prim(2.0, 10.0)
call multiply_add_prim(2.0, 2.0, 10.0)
Found expected exception:
Traceback (most recent call last):
File "/tmp/ipykernel_1796/2844449444.py", line 2, in <module>
square_add_prim(2., 10.)
File "/tmp/ipykernel_1796/3854395562.py", line 48, in func_wrapper
res = func(*args)
^^^^^^^^^^^
File "/tmp/ipykernel_1796/3275395289.py", line 17, in square_add_prim
return multiply_add_prim(a, a, b)
^^^^^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: Evaluation rule for 'multiply_add' not implemented
原始求值规则#
@trace("multiply_add_impl")
def multiply_add_impl(x, y, z):
"""Concrete implementation of the primitive.
This function does not need to be JAX traceable.
Args:
x, y, z: The concrete arguments of the primitive. Will only be called with
concrete values.
Returns:
the concrete result of the primitive.
"""
# Note: you can use the ordinary (non-JAX) NumPy, which is not JAX-traceable.
return np.add(np.multiply(x, y), z)
# Now, register the primal implementation with JAX:
multiply_add_p.def_impl(multiply_add_impl)
<function __main__.multiply_add_impl(x, y, z)>
assert square_add_prim(2., 10.) == 14.
call square_add_prim(2.0, 10.0)
call multiply_add_prim(2.0, 2.0, 10.0)
call multiply_add_impl(2.0, 2.0, 10.0)
|<- multiply_add_impl = 14.0
|<- multiply_add_prim = 14.0
|<- square_add_prim = 14.0
使用 jit 时会发生什么#
现在，如果您尝试使用 jit，您将收到一个 NotImplementedError。
with expectNotImplementedError():
api.jit(square_add_prim)(2., 10.)
call square_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>)
Found expected exception:
Traceback (most recent call last):
File "/tmp/ipykernel_1796/1813425700.py", line 2, in <module>
api.jit(square_add_prim)(2., 10.)
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/pjit.py", line 263, in cache_miss
executable, pgle_profiler, const_args) = _python_pjit_helper(
^^^^^^^^^^^^^^^^^^^^
NotImplementedError: Abstract evaluation for 'multiply_add' not implemented
抽象求值规则#
为了进行 JIT 编译以及其他变换，JAX 首先仅使用参数的形状和类型对其进行抽象求值。此抽象求值有多种用途：
获取计算中使用的 JAX 原语序列。此序列将被编译。
计算计算中使用的所有向量和操作的形状和类型。
例如，一个包含 3 个元素的向量的抽象可以是 ShapedArray(float32[3])，或者 ConcreteArray([1., 2., 3.])。在后一种情况下，JAX 使用实际的具体值作为抽象值包装。
from jax import core
@trace("multiply_add_abstract_eval")
def multiply_add_abstract_eval(xs, ys, zs):
"""Abstract evaluation of the primitive.
This function does not need to be JAX traceable. It will be invoked with
abstractions of the actual arguments
Args:
xs, ys, zs: Abstractions of the arguments.
Result:
a ShapedArray for the result of the primitive.
"""
assert xs.shape == ys.shape
assert xs.shape == zs.shape
return core.ShapedArray(xs.shape, xs.dtype)
# Now, register the abstract evaluation with JAX:
multiply_add_p.def_abstract_eval(multiply_add_abstract_eval)
<function __main__.multiply_add_abstract_eval(xs, ys, zs)>
如果您重新尝试应用 jit，您可以检查抽象求值是如何进行的，但您会收到另一个关于缺少实际 XLA 编译规则的错误。
with expectNotImplementedError():
api.jit(square_add_prim)(2., 10.)
call square_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
|<- square_add_prim = JitTracer<float32[]>
Found expected exception:
Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
app.launch_new_instance()
jax._src.source_info_util.JaxStackTraceBeforeTransformation: NotImplementedError: MLIR translation rule for primitive 'multiply_add' not found for platform cpu
The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
--------------------
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/tmp/ipykernel_1796/1813425700.py", line 2, in <module>
api.jit(square_add_prim)(2., 10.)
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/pjit.py", line 263, in cache_miss
executable, pgle_profiler, const_args) = _python_pjit_helper(
^^^^^^^^^^^^^^^^^^^^
NotImplementedError: MLIR translation rule for primitive 'multiply_add' not found for platform cpu
XLA 编译规则#
JAX 编译通过将每个原语编译为 XLA 操作图来工作。
这是向 JAX 添加新功能的最大障碍，因为 XLA 操作集是有限的，而 JAX 已经为其中大多数操作提供了预定义的原始。但是，XLA 包括一个 CustomCall 操作，可用于封装使用 C++ 定义的任意功能。
from jax._src.lib.mlir.dialects import hlo
@trace("multiply_add_lowering")
def multiply_add_lowering(ctx, xc, yc, zc):
"""The compilation to XLA of the primitive.
Given an mlir.ir.Value for each argument, return the mlir.ir.Values for
the results of the function.
Does not need to be a JAX-traceable function.
"""
return [hlo.AddOp(hlo.MulOp(xc, yc), zc).result]
# Now, register the lowering rule with JAX.
# For GPU, refer to the https://jax.net.cn/en/latest/Custom_Operation_for_GPUs.html
from jax.interpreters import mlir
mlir.register_lowering(multiply_add_p, multiply_add_lowering, platform='cpu')
您现在可以成功应用 jax.jit。请注意，JAX 首先抽象地求值函数，这会触发 multiply_add_abstract_eval 函数，然后编译它遇到的原语集，包括 multiply_add。此时 JAX 调用 multiply_add_lowering。
assert api.jit(lambda x, y: square_add_prim(x, y))(2., 10.) == 14.
call square_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
|<- square_add_prim = JitTracer<float32[]>
call multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x72c9f7618110>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x72c9f77f2840>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x72c9f77f27e0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x72c9f7604b70>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x72ca0fad2f60>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x72c9f7604bf0>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x72c9f77a61f0>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x72c9f7605550>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), platforms=None), Value(<block argument> of type 'tensor<f32>' at index: 0), Value(<block argument> of type 'tensor<f32>' at index: 1), Value(<block argument> of type 'tensor<f32>' at index: 2))
|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x72ca0d11d3b0>]
下面是 jit 的另一个用法，您仅针对第一个参数进行编译。请注意 square_add_prim 的第二个参数是具体的，这导致 multiply_add_abstract_eval 的第三个参数为 ConcreteArray。请注意，multiply_add_abstract_eval 可以与 ShapedArray 和 ConcreteArray 一起使用。
assert api.jit(lambda x, y: square_add_prim(x, y),
static_argnums=1)(2., 10.) == 14.
call square_add_prim(JitTracer<~float32[]>, 10.0)
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>, 10.0)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
|<- square_add_prim = JitTracer<float32[]>
call multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x72c9f76184d0>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x72c9f77f3380>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x72c9f77f3320>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x72c9f7606520>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x72ca0fad2f60>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x72c9f7606510>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x72c9f77a60d0>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x72c9f7606d80>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), platforms=None), Value(<block argument> of type 'tensor<f32>' at index: 0), Value(<block argument> of type 'tensor<f32>' at index: 1), Value(<block argument> of type 'tensor<f32>' at index: 2))
|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x72c9f7617870>]
前向自动微分#
JAX 以 Jacobian-Vector Product (JVP) 的形式实现前向自动微分（您可以在 高级自动微分 中了解更多）。
如果您尝试计算 jvp 函数，您会收到一个错误，因为您还没有告诉 JAX 如何对 multiply_add 原语进行微分。
# The second argument is set to `(2., 10.)` values where you
# evaluate the Jacobian, and the third argument `(1., 1.)`
# contains the values of the tangents for the arguments.
with expectNotImplementedError():
api.jvp(square_add_prim, (2., 10.), (1., 1.))
call square_add_prim(Traced<~float32[]>, Traced<~float32[]>)
call multiply_add_prim(Traced<~float32[]>, Traced<~float32[]>, Traced<~float32[]>)
Found expected exception:
Traceback (most recent call last):
File "/tmp/ipykernel_1796/459539105.py", line 5, in <module>
api.jvp(square_add_prim, (2., 10.), (1., 1.))
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/api.py", line 1890, in jvp
return _jvp(lu.wrap_init(fun, debug_info=debug_info("jvp", fun, primals, {})),
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: Differentiation rule for 'multiply_add' not implemented
from jax.interpreters import ad
@trace("multiply_add_value_and_jvp")
def multiply_add_value_and_jvp(arg_values, arg_tangents):
"""Evaluates the primal output and the tangents (Jacobian-vector product).
Given values of the arguments and perturbation of the arguments (tangents),
compute the output of the primitive and the perturbation of the output.
This method must be JAX-traceable. JAX may invoke it with abstract values
for the arguments and tangents.
Args:
arg_values: A tuple of arguments
arg_tangents: A tuple with the tangents of the arguments. The tuple has
the same length as the arg_values. Some of the tangents may also be the
special value `ad.Zero` to specify a zero tangent
Returns:
A pair of the primal output and the tangent.
"""
x, y, z = arg_values
xt, yt, zt = arg_tangents
_trace("Primal evaluation:")
# Now, you have a JAX-traceable computation of the output.
# Normally, you can use the multiply add (`ma`) primitive itself to compute the primal output.
primal_out = multiply_add_prim(x, y, z)
_trace("Tangent evaluation:")
# You must use a JAX-traceable way to compute the tangent. It turns out that
# the output tangent can be computed as (xt * y + x * yt + zt),
# which you can implement in a JAX-traceable way using the same "multiply_add_prim" primitive.
# You do need to deal specially with `Zero`. Here, you just turn it into a
# proper tensor of 0s (of the same shape as 'x').
# An alternative would be to check for `Zero` and perform algebraic
# simplification of the output tangent computation.
def make_zero(tan):
return lax.zeros_like_array(x) if type(tan) is ad.Zero else tan
output_tangent = multiply_add_prim(make_zero(xt), y, multiply_add_prim(x, make_zero(yt), make_zero(zt)))
return (primal_out, output_tangent)
# Register the forward differentiation rule with JAX:
ad.primitive_jvps[multiply_add_p] = multiply_add_value_and_jvp
# Tangent is: xt*y + x*yt + zt = 1.*2. + 2.*1. + 1. = 5.
assert api.jvp(square_add_prim, (2., 10.), (1., 1.)) == (14., 5.)
call square_add_prim(Traced<~float32[]>, Traced<~float32[]>)
call multiply_add_prim(Traced<~float32[]>, Traced<~float32[]>, Traced<~float32[]>)
call multiply_add_value_and_jvp((2.0, 2.0, 10.0), (1.0, 1.0, 1.0))
Primal evaluation:
call multiply_add_prim(2.0, 2.0, 10.0)
call multiply_add_impl(2.0, 2.0, 10.0)
|<- multiply_add_impl = 14.0
|<- multiply_add_prim = 14.0
Tangent evaluation:
call multiply_add_prim(2.0, 1.0, 1.0)
call multiply_add_impl(2.0, 1.0, 1.0)
|<- multiply_add_impl = 3.0
|<- multiply_add_prim = 3.0
call multiply_add_prim(1.0, 2.0, 3.0)
call multiply_add_impl(1.0, 2.0, 3.0)
|<- multiply_add_impl = 5.0
|<- multiply_add_prim = 5.0
|<- multiply_add_value_and_jvp = (14.0, 5.0)
|<- multiply_add_prim = Traced<float32[]>
|<- square_add_prim = Traced<float32[]>
前向自动微分的 JIT 编译#
您可以将 jit 应用于前向自动微分函数。
assert api.jit(lambda arg_values, arg_tangents:
api.jvp(square_add_prim, arg_values, arg_tangents))(
(2., 10.), (1., 1.)) == (14., 5.)
call square_add_prim(Traced<~float32[]>, Traced<~float32[]>)
call multiply_add_prim(Traced<~float32[]>, Traced<~float32[]>, Traced<~float32[]>)
call multiply_add_value_and_jvp((JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>), (JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>))
Primal evaluation:
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
Tangent evaluation:
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
|<- multiply_add_value_and_jvp = (JitTracer<float32[]>, JitTracer<float32[]>)
|<- multiply_add_prim = Traced<float32[]>
|<- square_add_prim = Traced<float32[]>
call multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x72c9f7657d10>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x72c9f7666430>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x72c9f76663d0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x72c9f76042d0>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x72ca0fad2f60>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x72c9f7605a30>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x72c9f7650c00>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x72c9f7678da0>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), platforms=None), Value(<block argument> of type 'tensor<f32>' at index: 0), Value(<block argument> of type 'tensor<f32>' at index: 1), Value(<block argument> of type 'tensor<f32>' at index: 2))
|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x72c9f7675cf0>]
call multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x72c9f7657d10>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x72c9f7666430>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x72c9f76663d0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x72c9f76042d0>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x72ca0fad2f60>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x72c9f7605a30>, all_default_mem_kind=True, lowering_cache={LoweringCacheKey(primitive=multiply_add, eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), effects=frozenset(), params=FrozenDict({}), platforms=('cpu',)): LoweringCacheValue(func=<jaxlib.mlir.dialects.func.FuncOp object at 0x72c9f7666580>, output_types=[RankedTensorType(tensor<f32>)], const_args=(), const_arg_avals=(), inline=True)}, cached_primitive_lowerings={}, traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x72c9f7650c00>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[])), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x72c9f7678f50>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), platforms=None), Value(<block argument> of type 'tensor<f32>' at index: 0), Value(<block argument> of type 'tensor<f32>' at index: 1), Value(<block argument> of type 'tensor<f32>' at index: 2))
|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x72c9f7676970>]
请注意，首先，您抽象地求值 multiply_add_value_and_jvp，这会抽象地求值原始值和切线值（共 3 次调用 ma 原语）。然后，您编译原语的 3 个出现。
反向自动微分#
如果您现在尝试使用反向自动微分，您会注意到 JAX 首先使用 multiply_add_value_and_jvp 来计算抽象值的前向自动微分，然后遇到一个 NotImplementedError。
在计算反向自动微分时，JAX 首先对前向自动微分代码 multiply_add_value_and_jvp 进行抽象求值，以获得计算输出切线的原语的追踪。
观察到 JAX 在进行抽象求值时，使用微分点的具体值和切线的抽象值。
请注意，JAX 使用特殊的抽象切线值
Zero来表示ma的第三个参数对应的切线。这反映了您不区分square_add_prim的第二个参数的事实，该参数流向multiply_add_prim的第三个参数。同时请注意，在切线的抽象求值过程中，您将
0.0作为第三个参数的切线值。这是因为在multiply_add_value_and_jvp的定义中使用了make_zero函数。
# This is reverse differentiation w.r.t. the first argument of `square_add_prim`
with expectNotImplementedError():
api.grad(square_add_prim)(2., 10.)
call square_add_prim(LinearizeTracer<~float32[]>, 10.0)
call multiply_add_prim(LinearizeTracer<~float32[]>, LinearizeTracer<~float32[]>, 10.0)
call multiply_add_value_and_jvp((TypedFloat(2.0, dtype=float32), TypedFloat(2.0, dtype=float32), 10.0), (Traced<~float32[]>, Traced<~float32[]>, Zero(~float32[])))
Primal evaluation:
call multiply_add_prim(TypedFloat(2.0, dtype=float32), TypedFloat(2.0, dtype=float32), 10.0)
call multiply_add_impl(TypedFloat(2.0, dtype=float32), TypedFloat(2.0, dtype=float32), 10.0)
|<- multiply_add_impl = 14.0
|<- multiply_add_prim = 14.0
Tangent evaluation:
call multiply_add_prim(TypedFloat(2.0, dtype=float32), Traced<~float32[]>, 0.0)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = Traced<float32[]>
call multiply_add_prim(Traced<~float32[]>, TypedFloat(2.0, dtype=float32), Traced<float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = Traced<float32[]>
|<- multiply_add_value_and_jvp = (14.0, Traced<float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
call multiply_add_abstract_eval(~float32[], ~float32[], float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = LinearizeTracer<float32[]>
|<- square_add_prim = LinearizeTracer<float32[]>
Found expected exception:
Traceback (most recent call last):
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/interpreters/ad.py", line 480, in get_primitive_transpose
return primitive_transposes[p]
~~~~~~~~~~~~~~~~~~~~^^^
KeyError: multiply_add
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
app.launch_new_instance()
jax._src.source_info_util.JaxStackTraceBeforeTransformation: NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'multiply_add' not implemented
The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
--------------------
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/tmp/ipykernel_1796/2155094905.py", line 3, in <module>
api.grad(square_add_prim)(2., 10.)
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/api.py", line 428, in grad_f
_, g = value_and_grad_f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'multiply_add' not implemented
上述错误是因为 JAX 在使用前向自动微分代码计算反向自动微分时缺少一个关键部分。
转置#
如前所述，在计算反向自动微分时，JAX 会获得一个使用前向自动微分计算切线的原语追踪。然后，JAX 会抽象地反向解释这个追踪，并为每个原语应用一个转置规则。
为了理解正在发生的事情，请考虑一个更简单的函数 f(x, y) = x * y + y 的示例。假设您需要在点 (2., 4.) 处进行微分。JAX 将从输入切线 xt 和 yt 的切线生成 ft 的 JVP 切线计算。
a = xt * 4.
b = 2. * yt
c = a + b
ft = c + yt
根据构建，切线计算始终是输入切线的线性运算。切线计算中可能出现的唯一非线性运算符是乘法，但此时其中一个操作数是常数。
JAX 将通过反向处理 JVP 计算来生成反向自动微分计算。对于切线计算中的每个操作，它将使用操作结果的共切线来累积操作所使用的变量的共切线。
# Initialize cotangents of inputs and intermediate variables:
xct = yct = act = bct = cct = 0.
# Initialize cotangent of the output:
fct = 1.
# Process `ft = c + yt`:
cct += fct
yct += fct
# Process `c = a + b`:
act += cct
bct += cct
# Process `b = 2. * yt`:
yct += 2. * bct
# Process `a = xt * 4.`:
xct += act * 4.
您可以验证此计算是否产生 xct = 4. 和 yct = 3.，它们是函数 f 的偏导数。
JAX 知道对于可能出现在 JVP 计算中的每个原语如何对其进行转置。概念上，如果原语 p(x, y, z) 对于常数 x 是关于参数 y 和 z 线性运算的，例如 p(x, y, z) = y*cy + z*cz，那么该原语的转置是：
p_transpose(out_ct, x, _, _) = (None, out_ct*cy, out_ct*cz)
请注意，p_transpose 接收原语输出的共切线以及对应于原语每个参数的值。对于线性参数，转置会获得未定义的 _ 值，对于其他参数，它会获得实际的常数。转置返回原语每个参数的共切线值，常数参数返回 None。
具体来说
add_transpose(out_ct, _, _) = (out_ct, out_ct)
mult_transpose(out_ct, x, _) = (None, x * out_ct)
mult_transpose(out_ct, _, y) = (out_ct * y, None)
@trace("multiply_add_transpose")
def multiply_add_transpose(ct, x, y, z):
"""Evaluates the transpose of a linear primitive.
This method is only used when computing the backward gradient following
`value_and_jvp`, and is only needed for primitives that are used in the JVP
calculation for some other primitive. You need a transposition for `multiply_add_prim`,
because you have used `multiply_add_prim` in the computation of the `output_tangent` in
`multiply_add_value_and_jvp`.
In this case, multiply_add is not a linear primitive. However, it is used linearly
w.r.t. tangents in `multiply_add_value_and_jvp`:
`output_tangent(xt, yt, zt) = multiply_add_prim(xt, y, multiply_add_prim(x, yt, zt))`.
Always one of the first two multiplicative arguments is a constant.
Args:
ct: The cotangent of the output of the primitive.
x, y, z: The values of the arguments. The arguments that are used linearly
get an ad.UndefinedPrimal value. The other arguments get a constant
value.
Returns:
A tuple with the cotangent of the inputs, with the value None
corresponding to the constant arguments.
"""
if not ad.is_undefined_primal(x):
# This use of multiply_add is with a constant "x".
assert ad.is_undefined_primal(y)
ct_y = ad.Zero(y.aval) if type(ct) is ad.Zero else multiply_add_prim(x, ct, lax.zeros_like_array(x))
res = None, ct_y, ct
else:
# This use of multiply_add is with a constant "y".
assert ad.is_undefined_primal(x)
ct_x = ad.Zero(x.aval) if type(ct) is ad.Zero else multiply_add_prim(ct, y, lax.zeros_like_array(y))
res = ct_x, None, ct
return res
ad.primitive_transposes[multiply_add_p] = multiply_add_transpose
现在您可以完成 grad 的运行。
assert api.grad(square_add_prim)(2., 10.) == 4.
call square_add_prim(LinearizeTracer<~float32[]>, 10.0)
call multiply_add_prim(LinearizeTracer<~float32[]>, LinearizeTracer<~float32[]>, 10.0)
call multiply_add_value_and_jvp((TypedFloat(2.0, dtype=float32), TypedFloat(2.0, dtype=float32), 10.0), (Traced<~float32[]>, Traced<~float32[]>, Zero(~float32[])))
Primal evaluation:
call multiply_add_prim(TypedFloat(2.0, dtype=float32), TypedFloat(2.0, dtype=float32), 10.0)
call multiply_add_impl(TypedFloat(2.0, dtype=float32), TypedFloat(2.0, dtype=float32), 10.0)
|<- multiply_add_impl = 14.0
|<- multiply_add_prim = 14.0
Tangent evaluation:
call multiply_add_prim(TypedFloat(2.0, dtype=float32), Traced<~float32[]>, 0.0)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = Traced<float32[]>
call multiply_add_prim(Traced<~float32[]>, TypedFloat(2.0, dtype=float32), Traced<float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = Traced<float32[]>
|<- multiply_add_value_and_jvp = (14.0, Traced<float32[]>)
|<- multiply_add_prim = LinearizeTracer<float32[]>
|<- square_add_prim = LinearizeTracer<float32[]>
call multiply_add_transpose(1.0, UndefinedPrimal(~float32[]), TypedFloat(2.0, dtype=float32), UndefinedPrimal(float32[]))
call multiply_add_prim(1.0, TypedFloat(2.0, dtype=float32), 0.0)
call multiply_add_impl(1.0, TypedFloat(2.0, dtype=float32), 0.0)
|<- multiply_add_impl = 2.0
|<- multiply_add_prim = 2.0
|<- multiply_add_transpose = (2.0, None, 1.0)
call multiply_add_transpose(1.0, TypedFloat(2.0, dtype=float32), UndefinedPrimal(~float32[]), 0.0)
call multiply_add_prim(TypedFloat(2.0, dtype=float32), 1.0, 0.0)
call multiply_add_impl(TypedFloat(2.0, dtype=float32), 1.0, 0.0)
|<- multiply_add_impl = 2.0
|<- multiply_add_prim = 2.0
|<- multiply_add_transpose = (None, 2.0, 1.0)
请注意两次调用 multiply_add_transpose。它们对应于 multiply_add_value_and_jvp 中 output_tangent 的两次使用 multiply_add_prim 的调用。第一次调用转置对应于最后一次使用 multiply_add_prim：multiply_add_prim(xt, y, ...)，其中 y 是常数 2.0。
反向自动微分的 JIT 编译#
请注意，multiply_add_value_and_jvp 的抽象求值仅使用抽象值。而在没有 JIT 的情况下，您使用了 ConcreteArray。
assert api.jit(api.grad(square_add_prim))(2., 10.) == 4.
call square_add_prim(LinearizeTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_prim(LinearizeTracer<~float32[]>, LinearizeTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_value_and_jvp((JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>), (Traced<~float32[]>, Traced<~float32[]>, Zero(~float32[])))
Primal evaluation:
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
Tangent evaluation:
call multiply_add_prim(JitTracer<~float32[]>, Traced<~float32[]>, JitTracer<~float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = Traced<float32[]>
call multiply_add_prim(Traced<~float32[]>, JitTracer<~float32[]>, Traced<float32[]>)
call multiply_add_abstract_eval(~float32[], ~float32[], float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = Traced<float32[]>
|<- multiply_add_value_and_jvp = (JitTracer<float32[]>, Traced<float32[]>)
|<- multiply_add_prim = LinearizeTracer<float32[]>
|<- square_add_prim = LinearizeTracer<float32[]>
call multiply_add_transpose(JitTracer<float32[]>, UndefinedPrimal(~float32[]), JitTracer<~float32[]>, UndefinedPrimal(float32[]))
call multiply_add_prim(JitTracer<float32[]>, JitTracer<~float32[]>, JitTracer<~float32[]>)
call multiply_add_abstract_eval(float32[], ~float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
|<- multiply_add_transpose = (JitTracer<float32[]>, None, JitTracer<float32[]>)
call multiply_add_transpose(JitTracer<float32[]>, JitTracer<~float32[]>, UndefinedPrimal(~float32[]), JitTracer<~float32[]>)
call multiply_add_prim(JitTracer<~float32[]>, JitTracer<float32[]>, JitTracer<~float32[]>)
call multiply_add_abstract_eval(~float32[], float32[], ~float32[])
|<- multiply_add_abstract_eval = float32[]
|<- multiply_add_prim = JitTracer<float32[]>
|<- multiply_add_transpose = (None, JitTracer<float32[]>, JitTracer<float32[]>)
call multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x72c9f76b2870>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x72c9f76bb4c0>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x72c9f77f0ee0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x72c9f76d1e30>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x72ca0fad2f60>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x72c9f76d1e20>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x72c9f76bd080>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[]), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x72c9f76d20c0>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), platforms=None), Value(<block argument> of type 'tensor<f32>' at index: 0), Value(<block argument> of type 'tensor<f32>' at index: 1), Value(<block argument> of type 'tensor<f32>' at index: 2))
|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x72c9f76de4b0>]
call multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x72c9f76b2870>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x72c9f76bb4c0>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x72c9f77f0ee0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x72c9f76d1e30>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x72ca0fad2f60>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x72c9f76d1e20>, all_default_mem_kind=True, lowering_cache={LoweringCacheKey(primitive=multiply_add, eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), avals_in=(ShapedArray(float32[]), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), effects=frozenset(), params=FrozenDict({}), platforms=('cpu',)): LoweringCacheValue(func=<jaxlib.mlir.dialects.func.FuncOp object at 0x72c9f76d8be0>, output_types=[RankedTensorType(tensor<f32>)], const_args=(), const_arg_avals=(), inline=True)}, cached_primitive_lowerings={}, traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x72c9f76bd080>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[]), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x72c9f76d22d0>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), platforms=None), Value(<block argument> of type 'tensor<f32>' at index: 0), Value(<block argument> of type 'tensor<f32>' at index: 1), Value(<block argument> of type 'tensor<f32>' at index: 2))
|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x72c9f76de8b0>]
批处理#
批处理变换将逐点计算转换为向量计算。如果您现在尝试，您将收到一个 NotImplementedError。
# The arguments are two vectors instead of two scalars.
with expectNotImplementedError():
api.vmap(square_add_prim, in_axes=0, out_axes=0)(np.array([2., 3.]),
np.array([10., 20.]))
call square_add_prim(Traced<float32[]>, Traced<float32[]>)
call multiply_add_prim(Traced<float32[]>, Traced<float32[]>, Traced<float32[]>)
Found expected exception:
Traceback (most recent call last):
File "/tmp/ipykernel_1796/1080163607.py", line 3, in <module>
api.vmap(square_add_prim, in_axes=0, out_axes=0)(np.array([2., 3.]),
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^
File "/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/api.py", line 1159, in vmap_f
out_flat = batching.batch(
^^^^^^^^^^^^^^^
NotImplementedError: Batching rule for 'multiply_add' not implemented
您需要指示 JAX 如何评估原始的批处理版本。在此特定情况下，multiply_add_prim 已经对任何维度的输入向量执行逐点操作，因此批处理版本可以使用相同的 multiply_add_prim 实现。
from jax.interpreters import batching
@trace("multiply_add_batch")
def multiply_add_batch(vector_arg_values, batch_axes):
"""Computes the batched version of the primitive.
This must be a JAX-traceable function.
Since the `multiply_add primitive` already operates point-wise on arbitrary
dimension tensors, to batch it you can use the primitive itself. This works as
long as both the inputs have the same dimensions and are batched along the
same axes. The result is batched along the axis that the inputs are batched.
Args:
vector_arg_values: A tuple of two arguments, each being a tensor of matching
shape.
batch_axes: The axes that are being batched. See vmap documentation.
Returns:
A tuple of the result, and the result axis that was batched.
"""
assert batch_axes[0] == batch_axes[1]
assert batch_axes[0] == batch_axes[2]
_trace("Using multiply_add to compute the batch:")
res = multiply_add_prim(*vector_arg_values)
return res, batch_axes[0]
batching.primitive_batchers[multiply_add_p] = multiply_add_batch
assert np.allclose(api.vmap(square_add_prim, in_axes=0, out_axes=0)(
np.array([2., 3.]),
np.array([10., 20.])),
[14., 29.])
call square_add_prim(Traced<float32[]>, Traced<float32[]>)
call multiply_add_prim(Traced<float32[]>, Traced<float32[]>, Traced<float32[]>)
call multiply_add_batch(([2. 3.], [2. 3.], [10. 20.]), (0, 0, 0))
Using multiply_add to compute the batch:
call multiply_add_prim([2. 3.], [2. 3.], [10. 20.])
call multiply_add_impl([2. 3.], [2. 3.], [10. 20.])
|<- multiply_add_impl = [14. 29.]
|<- multiply_add_prim = [14. 29.]
|<- multiply_add_batch = ([14. 29.], 0)
|<- multiply_add_prim = Traced<float32[]>
|<- square_add_prim = Traced<float32[]>
批处理的 JIT 编译#
下面是将 JIT 应用于批处理的示例。
assert np.allclose(api.jit(api.vmap(square_add_prim, in_axes=0, out_axes=0))
(np.array([2., 3.]),
np.array([10., 20.])),
[14., 29.])
call square_add_prim(Traced<float32[]>, Traced<float32[]>)
call multiply_add_prim(Traced<float32[]>, Traced<float32[]>, Traced<float32[]>)
call multiply_add_batch((JitTracer<float32[2]>, JitTracer<float32[2]>, JitTracer<float32[2]>), (0, 0, 0))
Using multiply_add to compute the batch:
call multiply_add_prim(JitTracer<float32[2]>, JitTracer<float32[2]>, JitTracer<float32[2]>)
call multiply_add_abstract_eval(float32[2], float32[2], float32[2])
|<- multiply_add_abstract_eval = float32[2]
|<- multiply_add_prim = JitTracer<float32[2]>
|<- multiply_add_batch = (JitTracer<float32[2]>, 0)
|<- multiply_add_prim = Traced<float32[]>
|<- square_add_prim = Traced<float32[]>
call multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x72c9f76b3110>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x72c9f76bbe20>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x72c9f76d8760>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x72c9f767b5a0>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x72ca0fad2f60>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x72c9f767adb0>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x72c9f76bd110>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[2]), ShapedArray(float32[2]), ShapedArray(float32[2])), avals_out=[ShapedArray(float32[2])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x72c9f76d3560>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), xla_metadata=None), platforms=None), Value(<block argument> of type 'tensor<2xf32>' at index: 0), Value(<block argument> of type 'tensor<2xf32>' at index: 1), Value(<block argument> of type 'tensor<2xf32>' at index: 2))
|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x72c9f76e88b0>]