patching shape query in TensorProxy to avoid prims

[jjsjann123]

We should avoid calling TensorProxy.shape directly during trace transformations, but use TensorProxy._shape instead.

The difference is that shape would leave a prims.shape symbol in the trace, if under a tracing context. Which is the case during grad transform.

Without this change, the trace will be corrupted with prims.shape calling on non-existing TensorProxy.

[mruberry]

This is really interesting. fyi @t-vi and @IvanYashchuk

As a note for the future, when we're transforming traces, maybe we shouldn't generally be operating within the trace context, to avoid accidentally putting operations into the trace. I guess if we did that then calling .shape on these tensors would throw an error that we're calling a tracing operation outside a trace context? If it doesn't, then maybe it should? Otherwise it would be easy for a developer to accidentally write .shape when they mean ._shape.

I also wonder if, in the future, we should make these queries more explicit. For example, instead of telling people to query ._shape to avoid appearing in a trace (which seems totally fine) we could have properties like .shape_notrace or a utility library like with functions like notrace_shape(x) to acquire metadata without appearing in a trace.

I don't think we have to pursue either of these ideas in this PR, of course.

[mruberry]

[jjsjann123]

err.. I should have made it clear that these three PRs are stacked. so they really should be merged in the order of #1673 #1693 #1696

[IvanYashchuk]

err.. I should have made it clear that these three PRs are stacked. so they really should be merged in the order of #1673 #1693 #1696

This happens 😅 I usually keep the stacked PRs in the draft mode to prevent merges