jax.typing
模块#
JAX 类型模块是 JAX 特定静态类型注释的存放地。此子模块正在开发中;要查看此处导出的类型的提案,请参阅 https://jax.ac.cn/en/latest/jep/12049-type-annotations.html。
当前可用的类型是
jax.Array
: 任何 JAX 数组或跟踪器的注释(即 JAX 转换中数组的表示)。jax.typing.ArrayLike
: 任何可以安全地隐式转换为 JAX 数组的值的注释;这包括jax.Array
、numpy.ndarray
,以及 Python 内置数字值(例如int
、float
等)和 numpy 标量值(例如numpy.int32
、numpy.float64
等)jax.typing.DTypeLike
: 任何可以转换为 JAX 兼容 dtype 的值的注释;这包括字符串(例如 ‘float32’、‘int32’)、标量类型(例如 float、np.float32)、dtypes(例如 np.dtype(‘float32’)),或具有 dtype 属性的对象(例如 jnp.float32、jnp.int32)。
我们可能会在未来的版本中添加其他类型。
JAX 类型最佳实践#
在公共 API 函数中为 JAX 数组添加注释时,我们建议使用 ArrayLike
作为数组输入,以及 Array
作为数组输出。
例如,您的函数可能如下所示
import numpy as np
import jax.numpy as jnp
from jax import Array
from jax.typing import ArrayLike
def my_function(x: ArrayLike) -> Array:
# Runtime type validation, Python 3.10 or newer:
if not isinstance(x, ArrayLike):
raise TypeError(f"Expected arraylike input; got {x}")
# Runtime type validation, any Python version:
if not (isinstance(x, (np.ndarray, Array)) or np.isscalar(x)):
raise TypeError(f"Expected arraylike input; got {x}")
# Convert input to jax.Array:
x_arr = jnp.asarray(x)
# ... do some computation; JAX functions will return Array types:
result = x_arr.sum(0) / x_arr.shape[0]
# return an Array
return result
JAX 的大多数公共 API 遵循此模式。请特别注意,我们建议 JAX 函数不要接受诸如 list
或 tuple
等序列来代替数组,因为这会导致 JAX 变换(如 jit()
)产生额外开销,并且在批次变换(如 vmap()
或 jax.pmap()
)中可能会产生意想不到的行为。有关此问题的更多信息,请参见 非数组输入 NumPy 与 JAX