Replace use_cuda with device parameter in AnnLoader#2360
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2360 +/- ##
==========================================
- Coverage 87.36% 85.35% -2.01%
==========================================
Files 49 49
Lines 7785 7790 +5
==========================================
- Hits 6801 6649 -152
- Misses 984 1141 +157
|
flying-sheep
left a comment
flying-sheep
left a comment
There was a problem hiding this comment.
thanks! Some small notes from me
| use_cuda: bool = _UNSET, | ||
| **kwargs, | ||
| ): | ||
| if use_cuda is not _UNSET: |
There was a problem hiding this comment.
we have a standard Empty.TOKEN object for this:
anndata/src/anndata/compat/__init__.py
Lines 101 to 102 in a0c4283
There was a problem hiding this comment.
Good call — switched to Empty.TOKEN in 6f01e74.
| shuffle: bool = False, | ||
| use_default_converter: bool = True, | ||
| use_cuda: bool = False, | ||
| device: str = "cpu", |
There was a problem hiding this comment.
please use a Literal type enumerating all valid values. (maybe imported from torch if it exists there)
There was a problem hiding this comment.
Done — used Literal["cpu", "cuda", "mps"] since torch doesn't export a device string literal. Fixed in 6f01e74.
|
Hi @flying-sheep, I've addressed both points (switched to Empty.TOKEN and used Literal type) in 6f01e74. Let me know if there's anything else! |
2 participants
Summary
use_cuda: boolwithdevice: str = "cpu"indefault_converterandAnnLoader.__init__, addressing the TODO comment in the code (# maybe replace use_cuda with explicit device option).use_cudais kept for backward compatibility but emits aFutureWarningwhen explicitly passed. Passing bothuse_cudaanddeviceraises aValueError.pin_memorylogic is now correctly guarded to only apply on CPU (it is a CPU-only operation).Motivation
The existing
use_cudaboolean is limiting — it only supports CPU or CUDA. The newdeviceparameter accepts any PyTorch device string (e.g.,"cpu","cuda","cuda:1","mps"), making the API more flexible and consistent with PyTorch conventions.Test plan
test_annloader_default_device— defaultdevice="cpu"produces CPU tensorstest_annloader_explicit_cpu_device— explicitdevice="cpu"workstest_annloader_use_cuda_deprecation_warning—use_cudaemitsFutureWarningtest_annloader_use_cuda_and_device_conflict— both params raisesValueError