Lightning-AI / Lightning-AI/lightning-thunder

Improve Conditional Statement Error Messages

Open
#1,247 1 comment 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

debugging enhancement
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

  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.

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.