Looking for a way to track python operations at node level

I am looking for a way to "do something" (or track results) after python executes something (even intermediate steps). I can track with frames executing line by line but I am missing some of the information on what is happening. I would rather have control of the execution of each node (maybe through the AST module).

An example:

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

In this line I want to do be able to do the following:

• Get the output shape and type of np.reshape(A, (R, Q, 1, P)) just after execution (left side of the binop)
• Get the output shape and type of the matmul @ operation (actual binop).
• Get the output shape and type of A after assignment.

Is this even possible?