pytorch / pytorch/pytorch

unable to export a model with bucketize

Open
#170,548 1 comment 0 reactions 0 assignees View on GitHub
export-triaged oncall: export oncall: pt2
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

```python
import torch

class Model(torch.nn.Module):
def forward(self, mask):
num_patches_per_side = 5
boundaries = torch.arange(
torch.tensor(1 / num_patches_per_side, dtype=torch.float32),
torch.tensor(1.0, dtype=torch.float32),
torch.tensor(1 / num_patches_per_side, dtype=torch.float32),
)
w_len = torch.tensor(1, dtype=boundaries.dtype) / mask.sum()
torch._check(w_len.item() > 0)
fractional_coords_w = torch.arange(
torch.tensor(0.0, dtype=boundaries.dtype),
torch.tensor(1 - 1e-6, dtype=boundaries.dtype),
w_len,
)
bucket_coords_h = torch.bucketize(fractional_coords_w, boundaries, right=True)
return bucket_coords_h

model = Model()
x = torch.zeros((10,), dtype=torch.bool)
x[::2] = True
expected = model(x)
ep = torch.export.export(model, (x,), dynamic_shapes=({0: torch.export.Dim.DYNAMIC},))
```

fails with the following error:

```
File "/home/xadupre/vv/this312/lib/python3.12/site-packages/torch/fx/experimental/symbolic_shapes.py", line 7664, in _evaluate_expr
raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression zuf0 > 0.0 (unhinted: zuf0 > 0.0). (Size-like symbols: none)

consider using data-dependent friendly APIs such as guard_or_false, guard_or_true and statically_known_true.
Caused by: (_refs/__init__.py:5247 in arange)
For more information, run with TORCH_LOGS="dynamic"
For extended logs when we create symbols, also add TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="zuf0"
If you suspect the guard was triggered from C++, add TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
For more debugging help, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit?usp=sharing

For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1

The following call raised this error:
File "/home/xadupre/github/experimental-experiment/_unittests/ut_torch_interpreter/test_onnx_export_aten.py", line 3139, in forward
fractional_coords_w = torch.arange(

The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.
```

### Versions

```
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Retpolines; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Retpoline
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected

Versions of relevant libraries:
[pip3] mypy_extensions==1.1.0
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.8.4.1
[pip3] nvidia-cuda-cupti-cu12==12.8.90
[pip3] nvidia-cuda-nvrtc-cu12==12.8.93
[pip3] nvidia-cuda-runtime-cu12==12.8.90
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cufft-cu12==11.3.3.83
[pip3] nvidia-curand-cu12==10.3.9.90
[pip3] nvidia-cusolver-cu12==11.7.3.90
[pip3] nvidia-cusparse-cu12==12.5.8.93
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-nccl-cu12==2.28.9
[pip3] nvidia-nvjitlink-cu12==12.8.93
[pip3] nvidia-nvtx-cu12==12.8.90
[pip3] onnx==1.21.0
[pip3] onnxruntime-gpu==1.24.0
[pip3] optree==0.17.0
[pip3] pytorch-triton==3.5.1+gitbfeb0668
[pip3] torch==2.11.0.dev20251215+cu128
[pip3] torchaudio==2.9.0
[pip3] torchvision==0.25.0.dev20251215+cu128
[pip3] triton==3.6.0+git8fedd49b
[conda] Could not collect
```

cc @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.