Python3.11下jaxtyping入门程序解析失败求助

问题背景 #

使用Python 3.11，搭配jaxtyping 0.2.36与typeguard 4.4.1编写类型注解时，出现AST解析失败的语法错误，核心报错为：

SyntaxError: invalid syntax

报错指向维度标注字符串（如"m n"）的解析环节。

相关依赖：

jaxtyping==0.2.36
numpy==2.1.3
torch==2.5.1
typeguard==4.4.1

测试代码：

from typeguard import typechecked
from jaxtyping import Float
from torch import Tensor


@typechecked
def matmul(a: Float[Tensor, "m n"], b: Float[Tensor, "n p"]) -> Float[Tensor, "m p"]:
    """矩阵乘法的类型注解示例"""
    raise NotImplementedError("暂未实现")

问题原因 #

jaxtyping 0.2.x系列版本仅兼容typeguard 2.x/3.x，而typeguard 4.x对类型注解的AST转换逻辑进行了重大调整，会将jaxtyping的维度标注字符串（如"m n"）当作Python表达式解析，导致语法错误。

解决方案 #

两种方案任选其一：

1. 降级typeguard到兼容版本
   执行以下命令安装兼容的typeguard版本：

   pip install typeguard==2.13.3
   

2. 升级jaxtyping到最新版本
   新版jaxtyping（0.3+）已适配typeguard 4.x，执行升级命令：

   pip install --upgrade jaxtyping
   

内容的提问来源于stack exchange，提问作者dspyz