Make GPU transforms more memory efficient

[bejaeger]

[bejaeger]

This change is part of the following stack:

• Make GPU transforms more memory efficient #887 ◀

_{Change managed by git-spice.}

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[gemini-code-assist]

Code Review

This pull request introduces memory-efficient chunked processing for TorchQuantileTransformer and TorchTruncatedSVD to bound peak memory usage during fitting and transformation. It also adds support for randomized SVD in TorchTruncatedSVD and implements a CPU fallback for the MPS backend to handle unsupported SVD operations. Review feedback suggests optimizing memory usage by replacing tensor allocations with Python scalars in torch.where calls and removing redundant tensor expansions.

[oscarkey]

if you're maxed out then totally fine to keep it like this, but did you consider using torch.vmap(chunk_size=blah) rather than implementing the chunking manually? I just checked and it seems to be supported on torch==2.5

[bejaeger]

Yep, using it in torch svd now to reduce line count! :)
For torch quantile I stick to the manual way because "PyTorch emits a 'performance drop' warning and falls back to a slower path".

[oscarkey]

great savings!!