wuyii8941 opened a new issue, #19689:
URL: https://github.com/apache/tvm/issues/19689
## Expected behavior
ONNX `AffineGrid` (opset 20) supports both 2D output (`size=[N,C,H,W]`,
theta shape `[N, 2, 3]`) and 3D output (`size=[N,C,D,H,W]`, theta shape `[N, 3,
4]`). PyTorch's `F.affine_grid` and ONNX Runtime both handle 3D.
## Actual behavior
The TVM converter raises:
```
ValueError: Only 2D AffineGrid (size=[N,C,H,W]) is supported
```
## Reproduction
```python
import numpy as np
import onnx
from onnx import helper, TensorProto, numpy_helper
import onnxruntime as ort
from tvm.relax.frontend.onnx import from_onnx
theta = helper.make_tensor_value_info("theta", TensorProto.FLOAT, [2, 3, 4])
Y = onnx.ValueInfoProto(); Y.name = "Y"
size = numpy_helper.from_array(np.array([2, 1, 3, 4, 5], dtype=np.int64),
"size")
node = helper.make_node("AffineGrid", ["theta", "size"], ["Y"],
align_corners=1)
g = helper.make_graph([node], "g", [theta], [Y], initializer=[size])
m = helper.make_model(g, opset_imports=[helper.make_opsetid("", 20)])
theta_v = np.random.randn(2, 3, 4).astype(np.float32)
print("ORT shape:", ort.InferenceSession(m.SerializeToString()).run(None,
{"theta": theta_v})[0].shape)
# ORT shape: (2, 3, 4, 5, 3)
inf = onnx.shape_inference.infer_shapes(m)
mod = from_onnx(inf) # ValueError: Only 2D AffineGrid is supported
```
## Root cause
`python/tvm/relax/frontend/onnx/onnx_frontend.py`, `AffineGrid._impl_v20`:
```python
# Only 2D is supported: size = [N, C, H, W]
if len(size_vals) != 4:
raise ValueError("Only 2D AffineGrid (size=[N,C,H,W]) is supported")
target_h, target_w = size_vals[2], size_vals[3]
grid = bb.emit(relax.op.image.affine_grid(theta, (target_h, target_w)))
return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 1]))
```
The underlying `relax.op.image.affine_grid` and `topi.image.affine_grid`
only target 2D. To support 3D, either:
- Extend `affine_grid` TOPI / Relax op to accept a 3D target shape
(preferable, matches PyTorch behavior), or
- At minimum, emit a `NotImplementedError` explicitly mentioning 3D as the
limitation (current `ValueError` reads as if 3D is invalid input rather than
missing implementation).
## Suggested fix
Add a 3D code path when `len(size_vals) == 5`. PyTorch's `F.affine_grid` for
3D produces a `[N, D, H, W, 3]` output. Implementation pattern:
```python
elif len(size_vals) == 5:
target_d, target_h, target_w = size_vals[2], size_vals[3], size_vals[4]
grid = bb.emit(relax.op.image.affine_grid_3d(theta, (target_d, target_h,
target_w)))
return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 4, 1]))
```
(requires adding a 3D `affine_grid` variant, or generalizing the existing
op).
## Impact
3D Spatial Transformer Networks (STNs), volumetric warping models, and any
opset-20+ ONNX export that uses 3D affine grids cannot be imported.
## Environment
- TVM: latest `main` (commit b172d5ea3)
- Python: 3.11
- ONNX Runtime: 1.24.4
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
