[ONNX] Implement torch.sym_sum and torch.sym_ite #170263
Conversation
Co-authored-by: justinchuby <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/170263
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 6 PendingAs of commit 4dd293b with merge base ca98449 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
release notes: onnx
justinchuby
added
module: onnx
|
@pytorchbot merge -i |
Merge startedYour change will be merged while ignoring the following 3 checks: pull / linux-jammy-cuda12.8-cudnn9-py3.10-clang12 / build, trunk / libtorch-linux-jammy-cuda12.8-py3.10-gcc11-debug / build, trunk / linux-jammy-cuda12.8-py3.10-gcc11-no-ops / build Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 3 jobs have failed, first few of them are: trunk / linux-jammy-rocm-py3.10 / test (default, 5, 6, linux.rocm.gpu.gfx942.1), trunk / linux-jammy-rocm-py3.10 / test (default, 4, 6, linux.rocm.gpu.gfx942.1), trunk / linux-jammy-rocm-py3.10 / test (default, 2, 6, linux.rocm.gpu.gfx942.1) Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
@pytorchbot merge -f "All related tests passed" |
|
The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Adds ONNX export support for `torch.sym_sum` and `torch.sym_ite` symbolic operations.
## Implementation
**torch/onnx/_internal/exporter/_torchlib/ops/symops.py**
- `sym_sum(args)`: N-ary addition using iterative ONNX Add operations. Handles empty sequences (returns 0) and single elements.
- `sym_ite(b, t, f)`: Conditional selection using ONNX Where operation.
**test/onnx/exporter/test_small_models_e2e.py**
- Test cases for both operations using dynamic shapes to trigger symbolic tracing.
Both implementations follow reference patterns from `torch/utils/_sympy/reference.py` and use the standard `@onnx_impl(trace_only=True)` decorator pattern.
## Example
```python
class Model(torch.nn.Module):
def forward(self, x):
# Sum of all dimensions
total = torch.sym_sum(x.shape)
# Conditional: max of first two dimensions
result = torch.sym_ite(x.shape[0] > x.shape[1], x.shape[0], x.shape[1])
return result
# Exports to ONNX with Add and Where ops
onnx_program = torch.onnx.export(Model(), torch.zeros(2, 3), dynamic_shapes=...)
```
Pull Request resolved: pytorch#170263
Approved by: https://github.com/titaiwangms
Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Adds ONNX export support for
torch.sym_sumandtorch.sym_itesymbolic operations.Implementation
torch/onnx/_internal/exporter/_torchlib/ops/symops.py
sym_sum(args): N-ary addition using iterative ONNX Add operations. Handles empty sequences (returns 0) and single elements.sym_ite(b, t, f): Conditional selection using ONNX Where operation.test/onnx/exporter/test_small_models_e2e.py
Both implementations follow reference patterns from
torch/utils/_sympy/reference.pyand use the standard@onnx_impl(trace_only=True)decorator pattern.Example
cc @justinchuby @titaiwangms