Feature Request: more general batched SVD operations?

[LionSR]

Hello Jax Developers and Community,

I've been delving into Jax's capabilities in Singular Value Decompositions (SVDs). While Jax provides support for gesvdjBatched (#1314), I've observed that this implementation comes with a matrix size restriction (m, n <= 32). On looking further into cuSOLVER's offerings, I stumbled upon gesvdaStridedBatched which doesn't seem to be a part of Jax currently.

Here are some considerations related to gesvdaStridedBatched:

1. General Matrix Support & No Size Restriction: Unlike gesvdjBatched, gesvdaStridedBatched offers support for general matrices without the strict size limitations, potentially offering more flexibility.

2. Usage in CuPy: I noticed that CuPy has integrated gesvdaStridedBatched into their framework. Here's their implementation. It's intriguing to see how they've leveraged it, and it made me wonder about its possible advantages if introduced in Jax.

3. Performance Advantages: An NVIDIA presentation (Page 14) touches upon the speed-ups achievable with Batched GESVDA. It could be of interest to evaluate if these performance metrics align with what Jax aims for.

On a related note, it seems that AMD's equivalent, ROCm, does not have such matrix size restrictions, as seen here. It does make me curious about the technical challenges or considerations behind these size restrictions.

I understand there are always multiple factors at play when deciding on integrations. It would be enlightening to hear the community's and developers' perspectives on this topic and whether gesvdaStridedBatched might find a place in Jax's roadmap.

Thank you for continually pushing the boundaries with Jax. Looking forward to learning more.

Edit 20230918: note that gesvdaStridedBatched can be called in the pytorch interface with the commit pytorch/pytorch@d136852 of the pull request pytorch/pytorch#74521.

Best wishes,
Ray

[adam-hartshorne]

This might be of interest,

Compositional Linear Algebra (CoLA)

CoLA is a framework for scalable linear algebra, automatically exploiting the structure often found in machine learning problems and beyond. CoLA supports both PyTorch and JAX.

https://github.com/wilson-labs/cola

[LionSR]

Pinging the author of #1314 @hawkinsp :-) happy to help if there are some instructions (#1099).

[TsingQAQ]

Hi JAX Community,

I wanted to follow up on this past discussion regarding SGESVDA (gesvdaStridedBatched) and check if it has/there are any plans for its integration into JAX. Given its ability to efficiently handle batched SVD without the size limitations of gesvdjBatched, it could be very useful.

One concrete motivation is randomized SVD, where gesvdaStridedBatched could significantly improve efficiency for low-rank approximations in large-scale settings (as one of the matrix dimensionality will typically be larger than 32). This is particularly important when working with large batch sizes, as current JAX implementations may face performance constraints.

If this is not currently a priority, are there any known workarounds in JAX? I checked the CoLA library (thanks to @adam-hartshorne for sharing!), but it doesn’t seem to support advanced batch SVD.

Would love to hear any insights or thoughts on this!

Best,

[dfm]

Thanks for checking in! I'd say that it's not a team priority to add support for this algorithm, but it would be welcomed as a PR! Probably the most useful commit to look at is 59e6e1c, which adds all the GPU kernels for SVD. That combined with the FFI tutorial might be enough to get started on a PR, and I'd be very happy to provide feedback on a PR thread!