jax.make_jaxpr

jax.make_jaxpr(fun: Callable, static_argnums: int | Iterable[int] = (), axis_env: Sequence[tuple[AxisName, int]] | None = None, return_shape: Literal[False] = False, abstracted_axes: Any | None = None) → Callable[..., core.ClosedJaxpr]

jax.make_jaxpr(fun: Callable, static_argnums: int | Iterable[int] = (), axis_env: Sequence[tuple[AxisName, int]] | None = None, return_shape: Literal[True] = False, abstracted_axes: Any | None = None) → Callable[..., tuple[core.ClosedJaxpr, Any]]

创建一个函数，该函数根据示例参数返回 fun 的 jaxpr。

参数:

• fun – 要计算其 jaxpr 的函数。其位置参数和返回值应为数组、标量或它们的标准 Python 容器（元组/列表/字典）。

• static_argnums – 请参阅 jax.jit() 的文档字符串。

• axis_env – 可选，一个对序列，其中第一个元素是轴名称，第二个元素是一个正整数，表示具有该名称的映射轴的大小。此参数在降低涉及并行通信集合的函数时很有用，它指定了由 jax.pmap() 的应用所设置的轴名称/大小环境。

• return_shape – 可选布尔值，默认为 False。如果为 True，则包装函数返回一个对，其中第一个元素是 fun 的 ClosedJaxpr 表示，第二个元素是一个 pytree，其结构与 fun 的输出相同，并且其叶子是具有 shape 和 dtype 属性的对象，这些属性表示输出叶子的相应类型。

返回:

一个 fun 的包装版本，当应用于示例参数时，返回这些参数上 fun 的 ClosedJaxpr 表示。如果参数 return_shape 为 True，则返回的函数改为返回一个对，其中第一个元素是 fun 的 ClosedJaxpr 表示，第二个元素是一个 pytree，表示 fun 输出的结构、形状、数据类型和命名形状。

jaxpr 是 JAX 用于程序跟踪的中间表示。jaxpr 语言基于带有 let-bindings 的简单类型一阶 lambda 演算。make_jaxpr() 使函数能够返回其 jaxpr，我们可以检查它以了解 JAX 内部的工作方式。返回的 jaxpr 是抽象到 ShapedArray 级别的 fun 的跟踪。内部还存在其他抽象级别。

我们在此不详细描述 jaxpr 语言的语义，而是提供一些示例。

>>> import jax
>>>
>>> def f(x): return jax.numpy.sin(jax.numpy.cos(x))
>>> print(f(3.0))
-0.83602
>>> jax.make_jaxpr(f)(3.0)
{ lambda ; a:f32[]. let b:f32[] = cos a; c:f32[] = sin b in (c,) }
>>> jax.make_jaxpr(jax.grad(f))(3.0)
{ lambda ; a:f32[]. let
    b:f32[] = cos a
    c:f32[] = sin a
    _:f32[] = sin b
    d:f32[] = cos b
    e:f32[] = mul 1.0:f32[] d
    f:f32[] = neg e
    g:f32[] = mul f c
  in (g,) }