`grid_sample` in Swift

We’re going to need something like affine_grid and grid_sample from PyTorch. (Which AFAICT is much the same as remap in OpenCV). Whilst we can simply use remap on CPU, that won’t help us on GPU. (Also remap doesn’t support batch transforms.)

It appears that tf doesn’t support this. Although, in classic tf fashion, some half-hearted implementations are dumped in various other places - most notably buried in Spatial Transformer in the official tf repo:

And something in tf.contrib:


And something in an external lib:

@saeta @clattner is this something that might be doable in Swift (and be performant on GPU)? Perhaps the s4tf team could see what the best approach might be?

Just discovered that cudnn already supports both these operations :slight_smile:


So how do we wrap a cudnn thing in s4tf most conveniently?

It seems the latest OpenCV also has cuda implementation of remap (in extra modules list)