Add grid sampling

[pxl-th]

This PR adds grid sampling for the 4D case.

This is needed for things like self-supervised depth estimation, where a neural network would learn a depth D for the current image and projection transformation P for the adjacent images and does "warping" of the images using D and P.
And this function does image sampling given projected coordinates (a.k.a. grid).

The implementation differs from the PyTorch's version in that it assumes align_corners = True and uses only bilinear interpolation. It was tested to give the same results as PyTorch's version.

PR with the CUDA kernels.

[pxl-th]

I've commented out failing test for the depthwise conv for now. Should I bring it back?

[CarloLucibello]

They were failing intermittently, mostly in the multi-threading environment. Since we have a few multithreaded methods, better have the multithreaded CI passing, so I'm ok with keeping those problematic tests out. We should open an issue for those depthwiseconv failures if there is no one already.

I'll try to review in the next few days, thanks for this contribution, at a first glance looks perfect already!

[ToucheSir]

@CarloLucibello RE CI failures, feel free to rename and re-purpose #359.

[CarloLucibello]

Suggested change

	Therefore, it should have values mostly in `[-1, 1]` range.
	Therefore, `x` and `y` should have values mostly in `[-1, 1]` range.

Also, why "mostly?"

[pxl-th]

Also, why "mostly?"

It was added as a hint that the values can be outside, but probably makes no sense.

[maxfreu]

Hi! Just to leave this here, as I ported the tri/bi/linear upsampling code from pytorch: 1) Setting align_corners = true is not the best choice, as the gradients change with changing input image size. Pytorch's default therefore is false. But that doesn't change any results much. 2) The code you wrote reminds me a lot of what I wrote for upsampling. The kernels probably can't be shared, although they look quite similar. But the utility functions can be shared I think - but that's maybe also for the future.

[gRox167]

Just to confirm, is there any plan to add 5D support for grid_sampling?
This is really helpful for doing 3D image registration, which is widely use in neuro and medical imaging field. Networks in medical image registration utilize this function extensively to wrap the moving image to a fixed image, for example VoxelMorph.

[CarloLucibello]

No active plans that I know of, but contribution are welcome. Maybe @maxfreu could provide some guidance.

[maxfreu]

Hi @gRox167 ! As imaging science phd student this is the perfect task for you :D In principle it should be quite simple:

1. Fork NNlib
2. Come up with simple test cases for the forward and backward pass that you want to pass and write those first.
3. Add specializations for AbstractArray{T,4} to the current implementation where it is needed.
4. Copy the code for the 4D case, change the specialization to 5D.
5. Add the depth dimension where sizes are checked etc.
6. Change the sampling code in the kernel. You can use the pytorch code as a reference, which you can find here for the GPU and here for the CPU.
7. Ping me or @pxl-th in case you hit bumps in the road :)
8. Enjoy how little code remains of the C++ hell.

It should be pretty straight forward and more of a writing than a thinking task. Optionally, you can simply use ChatGPT, show it the current julia and the pytorch code and tell it to expand it. I guess that should also work pretty well nowadays. The key to success in any case are good tests for the forward and backward pass.

[gRox167]

Hi @gRox167 ! As imaging science phd student this is the perfect task for you :D In principle it should be quite simple:

1. Fork NNlib

2. Come up with simple test cases for the forward and backward pass that you want to pass and write those first.

3. Add specializations for `AbstractArray{T,4}` to the current implementation where it is needed.

4. Copy the code for the 4D case, change the specialization to 5D.

5. Add the depth dimension where sizes are checked etc.

6. Change the sampling code in the kernel. You can use the pytorch code as a reference, which you can find [here](https://github.com/pytorch/pytorch/blob/f7000350905be5073892e0b23df681c0281be0f0/aten/src/ATen/native/cuda/GridSampler.cu#L156) for the GPU and [here](https://github.com/pytorch/pytorch/blob/f7000350905be5073892e0b23df681c0281be0f0/aten/src/ATen/native/GridSampler.cpp#L41) for the CPU.

7. Ping me or @pxl-th in case you hit bumps in the road :)

8. Enjoy how little code remains of the C++ hell.

It should be pretty straight forward and more of a writing than a thinking task. Optionally, you can simply use ChatGPT, show it the current julia and the pytorch code and tell it to expand it. I guess that should also work pretty well nowadays. The key to success in any case are good tests for the forward and backward pass.

Thanks for your guidance, I think I can take a try.

[gRox167]

Please refer to PR #627 .