Numba JIT Support + PyTorch JIT Support

[Skylion007]

Have you considered Numba JIT support to the Numpy methods in this library? With numba in particular, you could enable Parallelization and maybe even fastmath if the floating point errors are reasonable.

Also, several of the functions in the NUMBA tend give higher performance than the default Numpy ones. Main issue is that not all constructs/calls are supported so you sometimes have to reimplement it or change it to equivalent function signatures.

Also for the Torch examples, Torch.JIT support might give a slight performance boost, especially as the JIT graph optimizer improves. It also would allow it to be useful for lots of deployment.

You could write in a way so that jitting the code is not necessary and it would therefore not be an additional hard dependency.

[haruishi43]

@Skylion007 thanks for the recommendation! I would like to add numba supports if it gives higher performance. I'd like to create a benchmark and score each of the functions before this though. The loads are heavy in the grid_sampling functions, so I might start benchmarking there.

If you're versed in Numba or Pytorch's JIT, could you perhaps make a contribution? 😏

[Skylion007]

Your custom_grid sampling function in particular would benefit a lot since JIT can optimize most of those loops.

[haruishi43]

Thanks! I’ll look into it