Number proxies

[jjsjann123]

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Fixes # (issue).

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[t-vi]

What is our medium term plan w.r.t. defaul caching? If we need this for correctly handling #231 , it would seem that symbolic values should be the default but that in turn would mean that we want to have it work for our supported use cases.

[jjsjann123]

I'm hoping that #231 wouldn't need to the whole enable CACHE_OPTIONS.SYMBOLIC_VALUES thing. Looks like #231 has passed all CI, which feels promising.

The first step of this is to get number proxies to be plumbed through, I think that might be helpful for #231.
In terms of dynamic shape, we can slowly expanding its support.

Start off with allowing scalar input as number proxies, while still requiring tensor to be constant shape. I'm trying to figure out how/where to properly insert prologue_trace guard. Right now considering doing that from executor. i.e. nvfuser would require reduction dim(s) to be baked in as constant, while torchex doesn't care.

I'll try to come up with a design doc for review.

[jjsjann123]

Just adding anote to myself. nvfuserex trace looked a bit misleading:

import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast()
def computation(a, i0, i1):
  # a: "cuda:0 f32[8, 16, 32]"
  # i0: "int 0"
  # i1: "int 1"
  [t2] = nvFusion0(a, i0, i1)
    # t2 = prims.sum(a, (i0, i1))  # t2: "cuda:0 f32[32]"
  del a, i0, i1
  return t2

It looks like it's taking a dynamic reduction axis, but it's just baked in as constants.

Similarly, i0 and i1 are runtime value, we probably shouldn't have the trace indicate them as constant number in the trace where we have.

# i0: "int 0"
# i1: "int 1"

Similarly, output tensor t2 needs to show a symbolic shape.

[jjsjann123]

Note this one is still on-going.