cleanup affine grid image kernels

[pmeier]

Factoring out the changes to _apply_grid_transform from #7979. This pulls the squashing / unsquashing logic as well as the image.numel() == 0 handling into _apply_grid_transform to avoid repeating it in every kernel that uses it.

cc @vfdev-5

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/8004

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 12 New Failures, 10 Unrelated Failures

As of commit e9cad2d with merge base f96deba ():

NEW FAILURES - The following jobs have failed:

• linux (linux.12xlarge, cpu) / linux-job (gh)
• linux (linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh)
• unittests-linux (3.10, linux.12xlarge, cpu) / linux-job (gh)
• unittests-linux (3.11, linux.12xlarge, cpu) / linux-job (gh)
• unittests-linux (3.8, linux.12xlarge, cpu) / linux-job (gh)
• unittests-linux (3.8, linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh)
• unittests-linux (3.9, linux.12xlarge, cpu) / linux-job (gh)
• unittests-prototype (3.10, linux.12xlarge, cpu) / linux-job (gh)
• unittests-prototype (3.11, linux.12xlarge, cpu) / linux-job (gh)
• unittests-prototype (3.8, linux.12xlarge, cpu) / linux-job (gh)
• unittests-prototype (3.8, linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh)
• unittests-windows (3.10, windows.4xlarge, cpu) / windows-job (gh)

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• unittests-macos (3.10, macos-12) / macos-job (gh)
• unittests-macos (3.11, macos-12) / macos-job (gh)
• unittests-macos (3.8, macos-12) / macos-job (gh)
• unittests-macos (3.9, macos-12) / macos-job (gh)
• unittests-prototype (3.9, linux.12xlarge, cpu) / linux-job (gh)
• unittests-windows (3.11, windows.4xlarge, cpu) / windows-job (gh)
• unittests-windows (3.8, windows.4xlarge, cpu) / windows-job (gh)
• unittests-windows (3.9, windows.4xlarge, cpu) / windows-job (gh)
• windows (windows.4xlarge, cpu) / windows-job (gh)

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• unittests-macos (3.8, macos-m1-12) / macos-job (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[NicolasHug]

LGTM in principle, although I haven't double checked that all the new logic is strictly equivalent (hopefully it would be caught by the test). One Q below, but non-blocking as long as it's adressed / or if it's a non-issue

[NicolasHug]

Is it equivalent to cast to float16 before the call to _apply_grid_transform? It used to be done after.

[NicolasHug]

Also why is to(device) needed now?

[pmeier]

The float16 handling above is untouched and deals with the image. Here we are just aligning the displacement with the grid. This is the same as

vision/torchvision/transforms/v2/functional/_geometry.py

Lines 1794 to 1795 in f96deba

	if displacement.dtype != dtype or displacement.device != device:
	displacement = displacement.to(dtype=dtype, device=device)

but relying on the noop behavior of .to() to avoid the check on our side.