jax.eval_shape#

jax.eval_shape(fun, *args, **kwargs)[source]#

计算 fun 的形状/数据类型，无需任何浮点运算（FLOPs）。

此实用函数对于执行形状推断非常有用。其输入/输出行为由以下内容定义：

def eval_shape(fun, *args, **kwargs):
  out = fun(*args, **kwargs)
  shape_dtype_struct = lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype)
  return jax.tree_util.tree_map(shape_dtype_struct, out)

但是，它不直接应用 fun（这可能开销很大），而是利用 JAX 的抽象解释机制来评估形状，而无需执行任何浮点运算（FLOPs）。

使用 eval_shape() 也可以捕获形状错误，并将引发与评估 fun(*args, **kwargs) 相同的形状错误。

参数:

  • fun (可调用对象) – 应该评估其输出形状的函数。

  • *args – 一个位置参数元组，包含数组、标量或这些类型的（嵌套）标准 Python 容器（元组、列表、字典、命名元组，即 pytrees）。由于只访问 shape 和 dtype 属性，因此可以使用 jax.ShapeDtypeStruct 或其他鸭子类型为 ndarrays 的容器（但请注意，鸭子类型对象不能是命名元组，因为它们被视为标准 Python 容器）。

  • **kwargs – 一个关键字参数字典，包含数组、标量或这些类型的（嵌套）标准 Python 容器（pytrees）。与 args 类似，数组值只需鸭子类型为具有 shape 和 dtype 属性即可。

返回:

一个嵌套的 PyTree，其中叶子是 jax.ShapeDtypeStruct 对象。

返回类型:

输出

例如

>>> import jax
>>> import jax.numpy as jnp
>>>
>>> f = lambda A, x: jnp.tanh(jnp.dot(A, x))
>>> A = jax.ShapeDtypeStruct((2000, 3000), jnp.float32)
>>> x = jax.ShapeDtypeStruct((3000, 1000), jnp.float32)
>>> out = jax.eval_shape(f, A, x)  # no FLOPs performed
>>> print(out.shape)
(2000, 1000)
>>> print(out.dtype)
float32

所有通过 eval_shape() 传递的参数都将被视为动态参数；静态参数可以通过闭包包含，例如使用 functools.partial()。

>>> import jax
>>> from jax import lax
>>> from functools import partial
>>> import jax.numpy as jnp
>>>
>>> x = jax.ShapeDtypeStruct((1, 1, 28, 28), jnp.float32)
>>> kernel = jax.ShapeDtypeStruct((32, 1, 3, 3), jnp.float32)
>>>
>>> conv_same = partial(lax.conv_general_dilated, window_strides=(1, 1), padding="SAME")
>>> out = jax.eval_shape(conv_same, x, kernel)
>>> print(out.shape)
(1, 32, 28, 28)
>>> print(out.dtype)
float32