PyTorch's nn.functional.interpolate equivalent in JAX

[rolandgvc]

Hi guys, I was wondering if there's an existent equivalent to PyTorch's interpolate already in JAX (maybe .scipy?). I'm working on porting Torchvision's implementations of segmentation models to JAX/Flax. Let me know what I could use for this. Many thanks!
https://github.com/pytorch/vision/blob/master/torchvision/models/segmentation/_utils.py
https://pytorch.org/docs/stable/nn.functional.html

[shoyer]

Take a look at jax.image.resize and jax.scipy.ndimage. In particular, I think resize is exactly what you're looking for, though I'm not sure JAX has 100% of the options found in PyTorch.