search

jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.38k stars 2.79k forks source link

scipy.ndimage functions? #3647

Open benjaminpope opened 4 years ago

benjaminpope commented 4 years ago

Dear Jax team,

Before I bite the bullet and do this myself - I notice there's a stub for scipy.ndimage in Jax, but it doesn't have much in it.

I was thinking of using scipy.ndimage.rotate - is there a Jaxy way to do this or do we have to port all the scipy dependencies?

All the best,

Ben

shoyer commented 4 years ago

We would happily accept PRs to continue to fill out scipy.ndimage.

Most of the time, I expect the scipy functions will need to be rewritten to work well in JAX.

Vaibhavs10 commented 4 years ago

@samuela / @benjaminpope - Did you mean porting all the functions listed here: https://docs.scipy.org/doc/scipy/reference/ndimage.html over to jax.scipy.ndimage.py? I would be happy to help take this up, but would need some direction to get started.

benjaminpope commented 4 years ago

I mean, personally I was only going to use rotate and its dependencies, but I expect a lot of people will get mileage out of the ndimage functionality.

On Thu, 9 Jul 2020 at 6:14 pm, Vaibhav Srivastav notifications@github.com wrote:

@samuela https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_samuela&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=d2_F6fVP46oT__vbuTqmBNoqoE7E39iz-O2tK96KOaw&s=kdQguBogwamkORd4jbAZWQxZl1kLauXzNc4G-M7b6Mw&e= / @benjaminpope https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_benjaminpope&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=d2_F6fVP46oT__vbuTqmBNoqoE7E39iz-O2tK96KOaw&s=wEncRt9JQzw4SgNZGHlH9voDgZZxl3xkqw2vMlHt1Kg&e=

— You are receiving this because you were mentioned.

Reply to this email directly, view it on GitHub https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_google_jax_issues_3647-23issuecomment-2D655979606&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=d2_F6fVP46oT__vbuTqmBNoqoE7E39iz-O2tK96KOaw&s=qYxJrsnVQVOuH67G3zNYKEEK6ByGXMil1YzGEMqC5vw&e=, or unsubscribe https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_notifications_unsubscribe-2Dauth_ABN6YFP2PLKG2LWNOYW47DDR2V35RANCNFSM4OPPZTSA&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=d2_F6fVP46oT__vbuTqmBNoqoE7E39iz-O2tK96KOaw&s=KyXY3-NMVNP52wAntfX52Y5Vy-mY7nKssGz_vuPTehc&e= .

-- Dr Benjamin Pope NASA Sagan Fellow Center for Cosmology and Particle Physics // Center for Data Science New York University benjaminpope.github.io

Vaibhavs10 commented 4 years ago

@benjaminpope - That sounds good too. Can I help you out here and then work on porting the rest?

benjaminpope commented 4 years ago

Yeah! How can I help? The guts of Jax here are a little beyond me.

On Fri, Jul 10, 2020 at 2:12 AM Vaibhav Srivastav notifications@github.com wrote:

@benjaminpope https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_benjaminpope&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=_o7u3UxYj6C-EQsaWLhAg26CaQrMf_XOoidEAGfi0T8&s=spIa2mq6tGAE9GS2spxFqA4Ob827cWqPKG0PlaHwXis&e=

  • That sounds good too. Can I help you out here and then work on porting the rest?

— You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_google_jax_issues_3647-23issuecomment-2D656218726&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=_o7u3UxYj6C-EQsaWLhAg26CaQrMf_XOoidEAGfi0T8&s=6ZMQ9M8BkP4vQWFk65OORtPxJfSwVE3yyfAJLbhc2Hk&e=, or unsubscribe https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_notifications_unsubscribe-2Dauth_ABN6YFMV4OJATLMIFIUMYFDR2XUADANCNFSM4OPPZTSA&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=_o7u3UxYj6C-EQsaWLhAg26CaQrMf_XOoidEAGfi0T8&s=Q1CfCbtjDokGNhpce7iQfnB-nS9oC8g91mL7XauJ11E&e= .

-- Dr Benjamin Pope NASA Sagan Fellow Center for Cosmology and Particle Physics // Center for Data Science New York University benjaminpope.github.io

Vaibhavs10 commented 4 years ago

Sure. @jakevdp What do you think about this issue? I'd be willing to take a stab at it along with @benjaminpope. Do you have any recommendations on where we should start off with? I reckon we'll need to write a lot of low-level lax helper functions for this.

revannth commented 4 years ago

@jakevdp @benjaminpope @Vaibhavs10 I am willing to help out as well!

shoyer commented 4 years ago

Note that we just added jax.image.resize, which implements image resizing like PIL/TensorFlow: https://jax.readthedocs.io/en/latest/jax.image.html

So if there are cases where scipy.ndimage doesn't have quite the right API, we are willing to consider writing something new -- though we definitely would prefer to (re)implement established APIs, like those in scipy.ndimage.

As for the implementations itself, those should generally be done in terms of existing jax.lax or jax.numpy operations. With jit compilation, these often have similar performance to compiled alternatives. In cases where they don't, we can consider making use of XLA's CustomCall to call alternative implementations on CPU or GPU.

benjaminpope commented 4 years ago

OK cool. Just checking back in on this - have some cool research ideas to try if we have differentiable scipy.ndimage functions, but also probably a lot of work to do if nobody is planning to do this for other reasons.

On Sun, Jul 19, 2020 at 3:45 AM Stephan Hoyer notifications@github.com wrote:

Note that we just added jax.image.resize, which implements image resizing like PIL/TensorFlow: https://jax.readthedocs.io/en/latest/jax.image.html https://urldefense.proofpoint.com/v2/url?u=https-3A__jax.readthedocs.io_en_latest_jax.image.html&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=bbuUOkCzFl_ieXF5CGFDY3qbPoE1OFlzHqE31BmlU-E&s=s9RxV8pFv9TFWFDmjTmJyUUknEk2RVBN__DC3elo26I&e=

So if there are cases where scipy.ndimage doesn't have quite the right API, we are willing to consider writing something new -- though we definitely would prefer to (re)implement established APIs, like those in scipy.ndimage.

As for the implementations itself, those should generally be done in terms of existing jax.lax or jax.numpy operations. With jit compilation, these often have similar performance to compiled alternatives. In cases where they don't, we can consider making use of XLA's CustomCall to call alternative implementations on CPU or GPU.

— You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_google_jax_issues_3647-23issuecomment-2D660516754&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=bbuUOkCzFl_ieXF5CGFDY3qbPoE1OFlzHqE31BmlU-E&s=R56hq8AgRF8QpcIjswoGZFIVpX_1T2ypzApTKRsZIZ0&e=, or unsubscribe https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_notifications_unsubscribe-2Dauth_ABN6YFKS4QUQR5YV4FNFPLDR4HNQTANCNFSM4OPPZTSA&d=DwMCaQ&c=slrrB7dE8n7gBJbeO0g-IQ&r=_gjp43gRuTt_LjqJbH0jcZo0ePCz10Y4KegSkx-Ha0A&m=bbuUOkCzFl_ieXF5CGFDY3qbPoE1OFlzHqE31BmlU-E&s=6IcQSSUKl9ivddT5pwsxn-7dwx7wKnJJdTQVI9H91Lg&e= .

-- Dr Benjamin Pope NASA Sagan Fellow Center for Cosmology and Particle Physics // Center for Data Science New York University benjaminpope.github.io

SamPIngram commented 2 years ago

Interested in if you had any success in implementing the other ndimage functions? Or even if you gave it a go, but ran into issues? In a similar situation 2y on where I would like to use some of functions in Jax.

jakevdp commented 1 year ago

For anyone encountering this in a search: the dm_pix library implements several JAX-compatible implementations of image operations.