How do I rotate a 2D array?

[Logon27]

Hello, I am looking for a solution on how to rotate a 2D array by an arbitrary number of degrees (such as 45 degrees). I saw that google jax has scaling and translation, but no rotation (that I can find anyway)?
https://jax.readthedocs.io/en/latest/_autosummary/jax.image.scale_and_translate.html#jax.image.scale_and_translate

I am basically looking for an equivalent of scipy's ndimage rotate function.
https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.rotate.html

Does something similar to this exist for google jax?

[jakevdp]

The existing image manipulation tools within JAX can be found here:

• jax.image
• jax.scipy.ndimage

None of these contain image rotation functions, however. Note that there is a relevant feature request in #3647

As far as third party packages, it's worth checking out the dm_pix library, which is built on JAX and has a dm_pix.rotate function. I've not tried this, but it may be worth a look.

[Logon27]

Thank you for the response. I did see that issue, but seeing an unsolved issue from 3 years ago does not normally help that much. The dm_pix library is pretty much exactly what I am looking for. I never came across that library online when searching myself. I am slightly confused why there is this separate image processing library for jax developed by deepmind (also owned by google), but the functionality is not part of jax itself. It is pretty hard to find that dm_pix library unless you specifically search for the term 'google jax image processing'. Searching for individual image processing functions does not bring up results for this library for whatever reason.

[jakevdp]

In general, JAX's philosophy is to keep the main jax library small and focused on its core purpose. Lots of folks have built libraries on top of JAX for use in domain-specific settings; some of these libraries come from within Google and/or Deepmind, and some come from elsewhere.

Sorry you had trouble finding this other functionality.

[Logon27]

I appreciate you pointing me in the right direction. The philosophy part makes sense to some extent. However, Jax implements some image processing functionality (scaling and translation, but not rotation) and part of the scipy library (but not all functionality). Meanwhile these libraries implement some of the missing functionality that Jax does not have ported (both owned by Google). Which makes it slightly confusing, but I understand that the world is not perfect lol. Thanks again for your help!

[jakevdp]

Agreed it’s inconsistent - we’ve talked about removing the image-specific functionality from JAX but haven’t gone that far yet.