gemini-code-assist[bot] commented on code in PR #19863:
URL: https://github.com/apache/tvm/pull/19863#discussion_r3450283678
##########
python/tvm/topi/image/grid_sample.py:
##########
@@ -30,31 +30,38 @@ def affine_grid(data, target_shape):
Parameters
----------
data : tvm.Tensor
- 3-D with shape [batch, 2, 3]. The affine matrix.
+ 3-D with shape [batch, 2, 3] for 2D or [batch, 3, 4] for 3D. The
affine matrix.
- target_shape: list/tuple of two int
- Specifies the output shape (H, W).
+ target_shape: list/tuple of int
+ Specifies the output spatial shape (H, W) for 2D or (D, H, W) for 3D.
Returns
-------
Output : tvm.Tensor
- 4-D with shape [batch, 2, target_height, target_width]
+ [batch, 2, H, W] for 2D or [batch, 3, D, H, W] for 3D.
"""
assert target_shape is not None
- assert len(target_shape) == 2
- assert target_shape[0] > 1 and target_shape[1] > 1, (
- "target height/width should be greater than 1"
- )
+ assert len(target_shape) in (2, 3)
+ assert all(s > 1 for s in target_shape), "target spatial dims should be
greater than 1"
dtype = data.dtype
- y_step = tirx.const((2.0 - 1e-7) / (target_shape[0] - 1), dtype=dtype)
- x_step = tirx.const((2.0 - 1e-7) / (target_shape[1] - 1), dtype=dtype)
start = tirx.const(-1.0, dtype=dtype)
+ steps = [tirx.const((2.0 - 1e-7) / (s - 1), dtype=dtype) for s in
target_shape]
+
+ if len(target_shape) == 2:
+
+ def _compute(n, dim, i, j):
+ y = start + i * steps[0]
+ x = start + j * steps[1]
+ return data[n, dim, 0] * x + data[n, dim, 1] * y + data[n, dim, 2]
+
+ else:
- def _compute(n, dim, i, j):
- y = start + i * y_step
- x = start + j * x_step
- return data[n, dim, 0] * x + data[n, dim, 1] * y + data[n, dim, 2]
+ def _compute(n, dim, k, i, j):
+ z = start + k * steps[0]
+ y = start + i * steps[1]
+ x = start + j * steps[2]
+ return data[n, dim, 0] * x + data[n, dim, 1] * y + data[n, dim, 2]
* z + data[n, dim, 3]
Review Comment:
Instead of branching on the length of `target_shape` and duplicating the
computation logic for 2D and 3D, we can write a completely generic N-D affine
grid computation. This simplifies the code, improves maintainability, and
automatically supports other dimensions if needed in the future.
```python
def _compute(n, dim, *coords):
val = data[n, dim, len(target_shape)]
for r in range(len(target_shape)):
coord = start + coords[r] * steps[r]
val += data[n, dim, len(target_shape) - 1 - r] * coord
return val
```
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -3322,15 +3322,17 @@ def _impl_v20(cls, bb, inputs, attr, params):
else:
raise NotImplementedError(f"Dynamic size of type {type(size)} is
not supported")
- # 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]
-
- # Relax affine_grid outputs [N, 2, H, W]
- grid = bb.emit(relax.op.image.affine_grid(theta, (target_h, target_w)))
- # Permute to ONNX convention [N, H, W, 2]
- return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 1]))
+ if len(size_vals) == 4:
+ # 2D: size = [N, C, H, W]; relax affine_grid outputs [N, 2, H, W].
+ grid = bb.emit(relax.op.image.affine_grid(theta,
tuple(size_vals[2:])))
+ # Permute to ONNX convention [N, H, W, 2].
+ return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 1]))
+ if len(size_vals) == 5:
+ # 3D: size = [N, C, D, H, W]; relax affine_grid outputs [N, 3, D,
H, W].
+ grid = bb.emit(relax.op.image.affine_grid(theta,
tuple(size_vals[2:])))
+ # Permute to ONNX convention [N, D, H, W, 3].
+ return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 4, 1]))
+ raise ValueError("AffineGrid expects size to be [N,C,H,W] (2D) or
[N,C,D,H,W] (3D)")
Review Comment:
The implementation for 2D and 3D AffineGrid contains duplicated calls to
`bb.emit(relax.op.image.affine_grid(...))`. We can simplify this by determining
the permutation axes first, and then performing the common grid emission and
permutation steps once.
```python
if len(size_vals) == 4:
axes = [0, 2, 3, 1]
elif len(size_vals) == 5:
axes = [0, 2, 3, 4, 1]
else:
raise ValueError("AffineGrid expects size to be [N,C,H,W] (2D)
or [N,C,D,H,W] (3D)")
grid = bb.emit(relax.op.image.affine_grid(theta,
tuple(size_vals[2:])))
return bb.emit(relax.op.permute_dims(grid, axes=axes))
```
--
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]
