jax.make_jaxpr

内容

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][source]#
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]]

创建生成其 jaxpr 的函数，给定示例参数。

参数：

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

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

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

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

返回值：

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

jaxpr 是 JAX 用于程序跟踪的中间表示。 jaxpr 语言基于具有 let 绑定的简单类型的一阶 lambda 演算。make_jaxpr() 调整函数以返回其 jaxpr，我们可以检查它以了解 JAX 在内部执行的操作。返回的 jaxpr 是 fun 的跟踪，抽象到 ShapedArray 层次。其他抽象层次存在于内部。

我们在此处不详细描述 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 d
    f:f32[] = neg e
    g:f32[] = mul f c
  in (g,) }
目录