Lightning-AI / Lightning-AI/lightning-thunder
Support symbolic values in torch.arange and other tensor operations
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Thunder needs to support using symbolic scalar values as arguments to tensor creation operations, particularly `torch.arange` with symbolic start/end bounds.
## Current Behavior
Likely requires concrete values for tensor operations.
## Expected Behavior
Support symbolic values as arguments to:
- `torch.arange(symbolic_start, symbolic_end, ...)`
- `torch.zeros(symbolic_size, ...)`
- `.view(symbolic_shape, ...)` and other reshaping ops
- `.index_copy_(dim, symbolic_indices, ...)`
## Example from torch.compile
```py
# cache_position: "i64[1]" = torch.arange(cumulative_length, add, device=...)
cache_position: "i64[s50]" = torch.arange(
l_kwargs_past_key_values_layers_0_cumulative_length,
add, # add is Sym(s50 + s67)
device=device(type='cuda', index=0)
)
```
The output tensor has symbolic shape `"i64[s50]"`.
## Minimal Reproduction Case
```py
import torch
import thunder
@thunder.jit
def create_cache_position(cumulative_length: int, seq_len: int) -> torch.Tensor:
"""Create cache position indices using symbolic bounds."""
new_cumulative_length = cumulative_length + seq_len
# torch.arange with symbolic start and end
cache_position = torch.arange(
cumulative_length,
new_cumulative_length,
device='cuda'
)
return cache_position
result = create_cache_position(1024, 32)
assert result.shape == (32,)
assert result[0].item() == 1024
assert result[-1].item() == 1055
```
## Related Operations
- `torch.arange(symbolic, symbolic)`
- `torch.zeros(symbolic, *shape)`
- `tensor.view(-1, symbolic)`
- `tensor.scatter(1, symbolic_indices, value)`
## Success Criteria
- [ ] `torch.arange` accepts symbolic start/end
- [ ] Output tensor shape is correctly inferred as symbolic
- [ ] Other tensor ops support symbolic size arguments
- [ ] Generated code handles runtime dimensions correctly
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 Thunder/PyTorch minimal reproduction and compare its behavior with the listed symbolic operations. Trace how symbolic arguments and shapes are handled for torch.arange, torch.zeros, reshaping, and indexing, then verify the success criteria for runtime dimensions and symbolic output shapes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100