jaxtyping

<h1 align="center">jaxtyping</h1>

为以下内容提供类型注解和运行时类型检查：

1. JAX数组的形状和数据类型；（现在还支持PyTorch、NumPy和TensorFlow！）
2. PyTrees。

示例：

from jaxtyping import Array, Float, PyTree

# 接受具有匹配轴的浮点型2D数组
def matrix_multiply(x: Float[Array, "dim1 dim2"],
                    y: Float[Array, "dim2 dim3"]
                  ) -> Float[Array, "dim1 dim3"]:
    ...

def accepts_pytree_of_ints(x: PyTree[int]):
    ...

def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
    ...

安装

pip install jaxtyping

需要Python 3.9+。

JAX是一个可选依赖，仅用于一些特定于JAX的类型。如果没有安装JAX，这些类型将无法使用，但你仍然可以使用jaxtyping为PyTorch/NumPy/TensorFlow等提供形状/数据类型注解。

jaxtyping提供的注解与运行时类型检查包兼容，因此通常还会安装其中一个。最流行的两个是typeguard（对每个参数进行彻底检查）和beartype（随机检查参数的部分内容）。

文档

可在https://docs.kidger.site/jaxtyping查阅。

另请参阅：JAX生态系统中的其他库

总是有用的
Equinox：神经网络和核心JAX中尚未包含的所有内容！

深度学习
Optax：一阶梯度（SGD、Adam等）优化器。
Orbax：检查点（异步/多主机/多设备）。
Levanter：基础模型（如LLM）的可扩展+可靠训练。

科学计算
Diffrax：数值微分方程求解器。
Optimistix：根查找、最小化、不动点和最小二乘法。
Lineax：线性求解器。
BlackJAX：概率+贝叶斯采样。
sympy2jax：SymPy<->JAX转换；通过梯度下降训练符号表达式。
PySR：符号回归。（非JAX值得一提！）

Awesome JAX
Awesome JAX：更多JAX项目的长列表。