寻找一种在节点级别跟踪 python 操作的方法

Question

在 python 执行某些操作（甚至是中间步骤）之后，我正在寻找一种“做某事”（或跟踪结果）的方法。 我可以逐行跟踪执行的帧，但我错过了一些关于正在发生的事情的信息。 我宁愿控制每个节点的执行（也许通过 AST 模块）。

一个例子：

A = np.reshape(A, (R, Q, 1, P)) @ B

在这一行中，我希望能够做到以下几点：

• 执行后获取 output 形状和类型np.reshape(A, (R, Q, 1, P)) （binop的左侧）
• 获取 matmul @操作（实际 binop）的 output 形状和类型。
• 赋值后得到A的 output 形状和类型。

这甚至可能吗？