Can you change `use_cuda=False` to `device` so we can select `cuda`, `mps`, or `cpu`?

[jolespin]

Please describe your wishes and possible alternatives to achieve the desired result.

Looking forward to using anndata for more of my projects. Right now I'm testing a lot of my lighter weight pytorch models on apple silicon and use mps backend. Would it be possible for future versions to use device="auto" or similar to specify devices? If you have an auto feature the order could be something like cuda > mps > cpu or something.

[ilan-gold]

Hi @jolespin - could you clarify part of the codebase you are referring to?

[jolespin]

Apologies, I was referring to the PyTorch integration.

https://anndata.readthedocs.io/en/stable/generated/anndata.experimental.AnnLoader.html#anndata.experimental.AnnLoader

[ilan-gold]

@jolespin Since this would be a breaking change, would you be up for making it (especially since it wouldn't come out until our next minor release).

I think the way to go would be to deprecate the use_cuda argument, add a device argument, and implement as such in the meantime. What is the behavior with use_cuda=False? torch.tensor without usa_cuda keeps the data on the CPU?

[jolespin]

I think the easiest way to do it would be tensor.to(device) where device is either cpu, cuda, or mps with default being cpu. This would happen wherever there is a PyTorch object.

Use_cuda False keeps the data on CPU but if you have MPS (ie apple silicon) then you can't leverage those devices.