(l-onnx-doc-AffineGrid)=
# AffineGrid
(l-onnx-op-affinegrid-20)=
## AffineGrid - 20
### Version
- **name**: [AffineGrid (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#AffineGrid)
- **domain**: `main`
- **since_version**: `20`
- **function**: `True`
- **support_level**: `SupportType.COMMON`
- **shape inference**: `True`
This version of the operator has been available
**since version 20**.
### Summary
Generates a 2D or 3D flow field (sampling grid), given a batch of affine matrices theta
(https://pytorch.org/docs/stable/generated/torch.nn.functional.affine_grid.html).
An affine matrix `theta` is applied to a position tensor represented in its homogeneous expression. Here is an example in 3D:
```
[r00, r01, r02, t0] [x] [x']
[r10, r11, r12, t1] * [y] = [y']
[r20, r21, r22, t2] [z] [z']
[0, 0, 0, 1 ] [1] [1 ]
```
where `(x, y, z)` is the position in the original space, `(x', y', z')` is the position in the output space.
The last row is always `[0, 0, 0, 1]` and is not stored in the affine matrix. Therefore we have `theta` of shape `(N, 2, 3)` for 2D or `(N, 3, 4)` for 3D.
Input `size` is used to define grid of positions evenly spaced in the original 2D or 3D space, with dimensions ranging from `-1` to `1`.
The output `grid` contains positions in the output space.
When `align_corners=1`, consider `-1` and `1` to refer to the centers of the corner pixels (mark `v` in illustration).
```
v v v v
|-------------------|------------------|
-1 0 1
```
When `align_corners=0`, consider `-1` and `1` to refer to the outer edge of the corner pixels.
```
v v v v
|------------------|-------------------|
-1 0 1
```
#### Function Body
The function definition for this operator.
```
<
domain: "",
opset_import: ["" : 20]
>
AffineGrid (theta, size) => (grid)
{
one = Constant ()
two = Constant ()
zero = Constant ()
four = Constant ()
one_1d = Constant ()
zero_1d = Constant ()
minus_one = Constant ()
minus_one_f = CastLike (minus_one, theta)
zero_f = CastLike (zero, theta)
one_f = CastLike (one, theta)
two_f = CastLike (two, theta)
constant_align_corners = Constant ()
constant_align_corners_equal_zero = Equal (constant_align_corners, zero)
size_ndim = Size (size)
condition_is_2d = Equal (size_ndim, four)
N, C, D, H, W = If (condition_is_2d) ( N_then, C_then, D_then, H_then, W_then) {
N_then, C_then, H_then, W_then = Split (size)
D_then = Identity (one_1d)
}, else_branch: graph = g2 () => ( N_else, C_else, D_else, H_else, W_else) {
N_else, C_else, D_else, H_else, W_else = Split (size)
}>
size_NCDHW = Concat (N, C, D, H, W)
theta_3d = If (condition_is_2d) ( theta_then) {
gather_idx_6 = Constant ()
shape_23 = Constant ()
gather_idx_23 = Reshape (gather_idx_6, shape_23)
shape_N23 = Concat (N, shape_23)
gather_idx_N23 = Expand (gather_idx_23, shape_N23)
thetaN23 = GatherElements (theta, gather_idx_N23)
r1, r2 = Split (thetaN23)
r1_ = Squeeze (r1)
r2_ = Squeeze (r2)
r11, r12, t1 = Split (r1_)
r21, r22, t2 = Split (r2_)
r11_shape = Shape (r21)
float_zero_1d_ = ConstantOfShape (r11_shape)
float_zero_1d = CastLike (float_zero_1d_, theta)
float_one_1d = Add (float_zero_1d, one_f)
R1 = Concat (r11, r12, float_zero_1d, t1)
R2 = Concat (r21, r22, float_zero_1d, t2)
R3 = Concat (float_zero_1d, float_zero_1d, float_one_1d, float_zero_1d)
R1_ = Unsqueeze (R1, one_1d)
R2_ = Unsqueeze (R2, one_1d)
R3_ = Unsqueeze (R3, one_1d)
theta_then = Concat (R1_, R2_, R3_)
}, else_branch: graph = g4 () => ( theta_else) {
theta_else = Identity (theta)
}>
two_1d = Constant ()
three_1d = Constant ()
five_1d = Constant ()
constant_D_H_W_shape = Slice (size_NCDHW, two_1d, five_1d)
zeros_D_H_W_ = ConstantOfShape (constant_D_H_W_shape)
zeros_D_H_W = CastLike (zeros_D_H_W_, theta)
ones_D_H_W = Add (zeros_D_H_W, one_f)
D_float = CastLike (D, zero_f)
H_float = CastLike (H, zero_f)
W_float = CastLike (W, zero_f)
start_d, step_d, start_h, step_h, start_w, step_w = If (constant_align_corners_equal_zero) ( start_d_then, step_d_then, start_h_then, step_h_then, start_w_then, step_w_then) {
step_d_then = Div (two_f, D_float)
step_h_then = Div (two_f, H_float)
step_w_then = Div (two_f, W_float)
step_d_half = Div (step_d_then, two_f)
start_d_then = Add (minus_one_f, step_d_half)
step_h_half = Div (step_h_then, two_f)
start_h_then = Add (minus_one_f, step_h_half)
step_w_half = Div (step_w_then, two_f)
start_w_then = Add (minus_one_f, step_w_half)
}, else_branch: graph = h2 () => ( start_d_else, step_d_else, start_h_else, step_h_else, start_w_else, step_w_else) {
D_float_nimus_one = Sub (D_float, one_f)
H_float_nimus_one = Sub (H_float, one_f)
W_float_nimus_one = Sub (W_float, one_f)
D_equals_one = Equal (D, one)
step_d_else = If (D_equals_one) ( step_d_else_then) {
step_d_else_then = Identity (zero_f)
}, else_branch: graph = g6 () => ( step_d_else_else) {
step_d_else_else = Div (two_f, D_float_nimus_one)
}>
step_h_else = Div (two_f, H_float_nimus_one)
step_w_else = Div (two_f, W_float_nimus_one)
start_d_else = Identity (minus_one_f)
start_h_else = Identity (minus_one_f)
start_w_else = Identity (minus_one_f)
}>
grid_w_steps_int = Range (zero, W, one)
grid_w_steps_float = CastLike (grid_w_steps_int, step_w)
grid_w_steps = Mul (grid_w_steps_float, step_w)
grid_w_0 = Add (start_w, grid_w_steps)
grid_h_steps_int = Range (zero, H, one)
grid_h_steps_float = CastLike (grid_h_steps_int, step_h)
grid_h_steps = Mul (grid_h_steps_float, step_h)
grid_h_0 = Add (start_h, grid_h_steps)
grid_d_steps_int = Range (zero, D, one)
grid_d_steps_float = CastLike (grid_d_steps_int, step_d)
grid_d_steps = Mul (grid_d_steps_float, step_d)
grid_d_0 = Add (start_d, grid_d_steps)
zeros_H_W_D = Transpose (zeros_D_H_W)
grid_d_1 = Add (zeros_H_W_D, grid_d_0)
grid_d = Transpose (grid_d_1)
zeros_D_W_H = Transpose (zeros_D_H_W)
grid_h_1 = Add (zeros_D_W_H, grid_h_0)
grid_h = Transpose (grid_h_1)
grid_w = Add (grid_w_0, zeros_D_H_W)
grid_w_usqzed = Unsqueeze (grid_w, minus_one)
grid_h_usqzed = Unsqueeze (grid_h, minus_one)
grid_d_usqzed = Unsqueeze (grid_d, minus_one)
ones_D_H_W_usqzed = Unsqueeze (ones_D_H_W, minus_one)
original_grid = Concat (grid_w_usqzed, grid_h_usqzed, grid_d_usqzed, ones_D_H_W_usqzed)
constant_shape_DHW_4 = Constant ()
original_grid_DHW_4 = Reshape (original_grid, constant_shape_DHW_4)
original_grid_4_DHW_ = Transpose (original_grid_DHW_4)
original_grid_4_DHW = CastLike (original_grid_4_DHW_, theta_3d)
grid_N_3_DHW = MatMul (theta_3d, original_grid_4_DHW)
grid_N_DHW_3 = Transpose (grid_N_3_DHW)
N_D_H_W_3 = Concat (N, D, H, W, three_1d)
grid_3d_else_ = Reshape (grid_N_DHW_3, N_D_H_W_3)
grid_3d = CastLike (grid_3d_else_, theta_3d)
grid = If (condition_is_2d) ( grid_then) {
grid_squeezed = Squeeze (grid_3d, one_1d)
grid_then = Slice (grid_squeezed, zero_1d, two_1d, three_1d)
}, else_branch: graph = g2 () => ( grid_else) {
grid_else = Identity (grid_3d)
}>
}
```
### Attributes
* **align_corners - INT** (default is `'0'`):
if align_corners=1, consider -1 and 1 to refer to the centers of the corner pixels. if align_corners=0, consider -1 and 1 to refer to the outer edge the corner pixels.
### Inputs
- **theta** (heterogeneous) - **T1**:
input batch of affine matrices with shape (N, 2, 3) for 2D or (N, 3, 4) for 3D
- **size** (heterogeneous) - **T2**:
the target output image size (N, C, H, W) for 2D or (N, C, D, H, W) for 3D
### Outputs
- **grid** (heterogeneous) - **T1**:
output tensor of shape (N, H, W, 2) of 2D sample coordinates or (N, D, H, W, 3) of 3D sample coordinates.
### Type Constraints
* **T1** in ( `tensor(bfloat16)`, `tensor(double)`, `tensor(float)`, `tensor(float16)` ):
Constrain grid types to float tensors.
* **T2** in ( `tensor(int64)` ):
Constrain size's type to int64 tensors.