Lightning-AI / Lightning-AI/lightning-thunder

math.xxx calls in function on NumberProxy is not being traced.

Open
#526 1 comment 0 reactions 1 assignee View on GitHub

@jjsjann123 is already working on this.

Since Jun 10, 2024.

bug dynamic constraints symbolic values
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🐛 Bug

As per title suggests, with symbolic values cache, NumberProxy operation via math.xxx calls are not traced.

### To Reproduce

```
import torch
import thunder
import math

def foo(a, b):
return math.fmod(a, b)

a = 1.2
b = 2.0

jfoo = thunder.jit(foo, cache="symbolic values")
out = jfoo(a, b)

print("\n\tcompute last trace:\n", thunder.last_traces(jfoo)[-1])
```

Looking at the trace below, we do not have `prims.fmod` in trace but it just baked in the result and this is wrong.
```
compute last trace:
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation():
return 1.2
```

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.