Dyanmic Shape TensorProxy

[jjsjann123]

🚀 Feature

This is a continuation of the work in #262 . The phase 2 proposal targets the re-use of compiled thunder program by supporting TensorProxy with dynamic shape and adjusting caching policy accordingly.

This is an umbrella issue to track progress. I would like to leave detailed technical discussion in separate issues and link them. This section will be gradually populated as we start picking up speed:

• shape propagation modeling in thunder

  • [ ] Enabling NumberProxy in TensorProxy.shape (toy PR #1027)
  • [x] Adding shape prim (PR #1113)
  • [x] Issue that blocks the #1182 . Should be resolved in #1201
  • [ ] shape logic in prim needs to be modeled in trace Issue: #471
  • [ ] modeling shape queries: #1253 . Will need two PR to resolve, PR 1 #1260

• Static constraints and prologue checks

  • [ ] #1133

• Handling of existing shape checks

[jjsjann123]

Linking feature request #581

I think a better design should make this a constraint, rather than opaquely treat it as a cache miss by backend. Will target this when we proceed to static constraints.

[jjsjann123]

Linking issue reported by @t-vi https://github.com/Lightning-AI/lightning-thunder/issues/1257.