Lightning-AI / Lightning-AI/lightning-thunder
Improve Conditional Statement Error Messages
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature: Improve Conditional Statement Error Messages in Thunder JIT
### Motivation
When attempting to use conditional statements (e.g., if a.sum() > 0:) in Thunder JIT functions, the resulting error message is somewhat informative but can be enhanced. The message currently states that bool is not implemented, but it could go further in indicating that conditional statements are generally not supported in Thunder JIT.
The current error message can leave users uncertain about the exact cause of the issue and what alternatives are available. Improving this error feedback will help developers troubleshoot more effectively and understand the limitations of Thunder JIT more clearly.
Here is the repro script:
```py
import thunder
import torch
def foo(a, b):
if a.sum() > 0:
return b
return a + b
jfoo = thunder.jit(foo)
a = torch.full((2, 2), 1)
b = torch.full((2, 2), 3)
result = jfoo(a, b)
```
Which results in the following error (I redacted most of lines for simplicity):
```log
Traceback (most recent call last):
(...) line 5, in foo
if a.sum() > 0:
File "/home/mmikulski/miniconda3/envs/report2/lib/python3.10/site-packages/thunder/core/interpreter.py", line 1343, in impl
return dunder_bool(x)
File "/home/mmikulski/miniconda3/envs/report2/lib/python3.10/site-packages/thunder/core/interpreter.py", line 1253, in wrapping_wrapper
res = ufn(*uargs, **ukwargs)
File "/home/mmikulski/miniconda3/envs/report2/lib/python3.10/site-packages/thunder/core/proxies.py", line 1304, in __bool__
raise NotImplementedError
NotImplementedError
```
### Pitch
* Explicitly mention that conditional statements are not supported in Thunder JIT.
* Provide suggestions or alternatives for users, such as restructuring logic to avoid direct conditional statements.
* Optionally link to relevant documentation or examples that illustrate supported patterns for similar use cases.
### Alternatives
* Implementing limited conditional logic support within Thunder JIT, if feasible, to handle simple use cases.
### Additional context
The [Sharp Edges](https://lightning-thunder.readthedocs.io/en/latest/basic/sharp_edges.html) section of Thunder Documentation could contain information about conditional statements. Ideally, the error message would lead there.
cc @carmocca @apaz-cli
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by running the provided repro script, then read thunder/core/interpreter.py around impl and thunder/core/proxies.py around __bool__. Review the Sharp Edges documentation section as well; done means the conditional-statement failure clearly says this pattern is unsupported and points users toward documented alternatives.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, documentation
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100