Lightning-AI / Lightning-AI/lightning-thunder

Support symbolic values in torch.arange and other tensor operations

Open
#2,738 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement symbolic values
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.