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.52k stars 2.8k forks source link

`scipy.sparse.csgraph.connected_components` implementation #24737

Open vboussange opened 1 week ago

vboussange commented 1 week ago

Hey there, It would be awesome to have scipy.sparse.csgraph.connected_components implemented in JAX. I guess this relates to an old request (https://github.com/jax-ml/jax/discussions/9214) that I am reviving. We could get inspiration from the GPU-friendly cugraph implementation. Is this something that could be considered? Cheers!

jakevdp commented 1 week ago

Hi – this would probably not be a great fit for the JAX package itself (see the discussion in JEP 18137) but we'd welcome efforts in companion projects to implement this kind of functionality.

vboussange commented 6 days ago

I see, thanks! I guess that this would nicely fit in a JAX-based Graph library (Graphs.jl or NetworkX equivalent). Is this a companion project that has already been discussed?