ENH: add a float32-only device

[ev-br]

Add a new device which only supports single precision floats and does not support double precision floats.

The new device mimics torch "mps" device in that it does not have f64 but supports int64---unlike JAX which either has both or none.

• Make inspection APIs (dtypes, default_dtypes) device-aware;
• Make creation functions (ones, empty etc) use the device-specific default dtype when given dtype=None, device=f32_device

TODO:

• Finish adding device-specific default dtypes to all creation functions
• device-specific default dtypes for fft.{fftfreq, rfftfreq}
• Audit the codebase for missing device= arguments in internal constructions

TBD:

• re-purpose an existing device2 or add a new F32_device (if so, bikeshed the name)
• standardize the behavior for incompatible arguments: dtype=float64, device=f32_only_device should raise? torch "mps" tensors raise a TypeError, follow it or mandate a ValueError?

Intends to close gh-64,
Gives a way to close gh-38 --- if we have a f32-only device, we probably do not need a global flag
Addresses a large part of gh-70

Cross-ref the spec RFC to allow for missing dtypes , data-apis/array-api#998 --- note that this array-api-strict PR can only land after the spec is updated;
Also cross-ref the test suite tracker data-apis/array-api-tests#431: the test suite is actually fairly far along.

EDIT: data-apis/array-api#1005 is I believe the spec update to make this spec-compliant.

[betatim]

I had a quick look at the diff, I think this is going in a direction I like.

One thought to which I have no good answer: would it make sense to repurpose device2 as the one that doesn't support float64 instead of adding "yet another device"? A reason to reuse it is to avoid having many devices, a reason to have a new device is "explicit is better than implicit" (and it is basically free to add devices).

Interested what other think about this.

[ogrisel]

I am fine with repurposing device2.

[ogrisel]

I think it would make sense for the float32 only device to raise an exception on __setitem__ operations that trigger an implicit data transfer.

Note that this is not the case for PyTorch on the MPS devices for some reason:

>>> import torch
>>> a = torch.ones(10, device="mps")
>>> b = 2 * torch.ones(10)  # CPU
>>> a + b
Traceback (most recent call last):
  Cell In[9], line 1
    a + b
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, mps:0 and cpu!
>>> b[:5] = a[:5]  # implicit MPS to CPU transfer does not raise!
>>> b
tensor([1., 1., 1., 1., 1., 2., 2., 2., 2., 2.])

I would have expected the last statement to also raise RuntimeError for consistency. I think array-api-strict should raise in this case.

EDIT: actually, this is not related to dtypes. Maybe I should open a dedicated issue for this.

EDIT2: done: #207

[ev-br]

Okay, this is ready from my side.

Re yet another device vs repurpose device2, the latter is more churn; the former works today and a new device is free, so am opting to adding a new device unless someone has a strong opinion to the contrary.

[ev-br]

Rebased on top of gh-212.
@betatim could I tempt you to take a look at this PR if you've the time?

[rgommers]

This indeed seems useful, thanks for pushing this forward!

The new device mimics torch "mps" device in that it does not have f64 but supports int64---unlike JAX which either has both or none.

Also MLX on GPU: https://ml-explore.github.io/mlx/build/html/python/data_types.html

[betatim]

I like having this device! Here some review comments.

We need some kind of nice name for this device

[ev-br]

Thanks for the review @betatim ! I pushed an update which addressed most of your comments. Also took liberty to resolve those GH comments where an update seemed to just go with your suggestions, so that they don't distract from the few comments which are less clear-cut and where I'd be looking for more input.

[betatim]

Thanks for working on this and the rounds of nice review!

LGTM

[betatim]

@rgommers do you want to look at this or ponder things about it? Otherwise hit "merge" after reading this.

I will come back tomorrow or the day after to merge this if it is not yet merged.

[rgommers]

Otherwise hit "merge" after reading this.

✅ 🚀