ENH: add a float32-only device by ev-br · Pull Request #206 · data-apis/array-api-strict

ev-br · 2026-04-24T12:47:15Z

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 · 2026-04-27T07:06:08Z

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 · 2026-04-30T07:46:47Z

I am fine with repurposing device2.

ogrisel · 2026-04-30T07:52:17Z

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 · 2026-05-25T11:33:08Z

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.

As a part of a rebase/conflict resolution, this commit simply moves the dlpack/device updates to _devices.py. Co-authored-by: Tim Head <betatim@gmail.com>

ev-br · 2026-06-30T15:55:07Z

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

rgommers · 2026-07-06T10:29:04Z

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

betatim · 2026-07-06T12:11:49Z

+    _complex_floating_dtypes, _numeric_dtypes
+)
+
+_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "F32_device")


How about spelling out the name of the device a bit more?

My first thought was "float32_only_device" but I think that is confusing because you might think it means the only type available is float32. So maybe no_float64_device is more obvious. I do wonder if we could call it just no_float64, the device postfix feels a bit superfluous, but it does make the string meaningful by itself.

Suggested change

_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "F32_device")

_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "no_float64_device")

Thanks! Since the vast majority of usage is xp.Device("no_floa64") with the word Device already used, I went with no _device" postfix in the updated PR. If xp.Device("no_float64_device")` reads better to you @betatim , I'll change to that.

[a totally unneccessary explanation: English is a second language for me, and my first language has a rather strong preference for not repeating words even at a cost of breaking the word order in a sentence, thus the choice above; all that said, I'm not going to argue the point if you prefer Device("no_float64_device")]

betatim · 2026-07-06T12:50:48Z

+    if dtype is None and device is not None:
+        res_dtype = DType(res.dtype)
+        if not device_supports_dtype(device, res_dtype):
+            # find out the default dtype for the device


I think this does the right thing, but I was surprised by it. I was expecting the code to find the default dtype for the device, given that no dtype was passed to asarray, and then create the Numpy array with that information. Instead we first create the Numpy array, then check if the dtype it selected is supported by the device. This feels "backwards". I think most of that is because in my head I'm thinking "whether a device is supported or not, and if it is the default device or not are two different questions". As in, on some fictitious device the default dtype chosen by Numpy is supported, but not the default. Then this code would no longer do the right thing.

I think the reason the code is the way it is, is that without the help of np.array we'd have to work out which group of dtypes the input belongs to. For example in the case of array_api_strict.asarray([1,2,3.141], device=foobar) we'd have to determine that we need the floating point dtypes and for array_api_strict.asarray([1,2,3], device=foobar) we need the integer dtypes.

Maybe we can get the best of both worlds by keeping the current code and adding a check that the selected res_dtype is a allowed default dtype for device.

I think the below comment represents the current thinking and explains a bit what is happening here. I'd probably add this no matter what (unless we completely rewrite all this):

Suggested change

# find out the default dtype for the device

# The dtype selected by Numpy might not be the default dtype

# on this device. If the dtype is not supported by the device we

# try to find one that is, aka the default dtype of this device.

Agreed that the code is a bit awkward. And indeed I'd rather lean on numpy and correct afterwards and avoid scanning the inputs to find out the dtype from scratch which can be bit of a rabbit hole, too.

As in, on some fictitious device the default dtype chosen by Numpy is supported, but not the default. Then this code would no longer do the right thing.

So that would be a pytorch "cpu" device with the default pytorch settings, which supports float64 but defaults to float32. That said, I don't know if we need to go all the way here?

For now, I kept the adjust-the-dtype dance and added your explanatory comment, hoping that that'd be a reasonable compromise?

betatim · 2026-07-06T13:01:10Z

+    return full_like(a, fill_value=42.0, *args, **kwds)
+
+
+class TestDefaultDType:


Did you pick the class based testing approach on purpose? I often have AI do this, even if all of the rest of the project uses test functions. Unless there is a good reason for this I'd vote for keeping it consistent and using test functions. Not classes

It's just my default mode (that it matches what AI does is by accident, I swear I'm made of meat and bone :-)) to group related tests into a class, since I find it easier to navigate a test module with a dozen classes with a dozen test functions each, than a flat list of 144 test functions. Here for example, TestDefaultDType::test_eye clearly tests the eye function for its default dtype behavior, as opposed to a module-level test_eye which tests the eye function yes, but I need to read the test to find out which behaviors of this function are being tested.

So whether it's on purpose is yes and no: no I did not give it much thought, yes I've a mild preference towards this style; the preference is mild indeed, so if this irks the reviewer, I'm happy to remove the test class. Yay/Nay?

Co-authored-by: Tim Head <betatim@gmail.com>

ev-br · 2026-07-06T16:09:53Z

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 than clear-cut and I'd be looking for more input.

This was referenced Apr 24, 2026

Make one of the devices have a lower maximum float precision #64

Open

Improvements to device support #70

Open

ev-br force-pushed the f32_device branch 4 times, most recently from 5cd3619 to 6ba0f5b Compare April 25, 2026 09:39

ev-br mentioned this pull request Apr 26, 2026

[RFC]: Updates for float16/bfloat16 and for dtypes that are lacking full support in libraries data-apis/array-api#998

Open

ev-br force-pushed the f32_device branch from 6583851 to 5e72e78 Compare April 26, 2026 18:56

betatim reviewed Apr 27, 2026

View reviewed changes

Comment thread array_api_strict/_info.py

ev-br mentioned this pull request May 2, 2026

Raise on implicit device transfer in __setitem__ #207

Open

ev-br added this to the 2.6 milestone May 9, 2026

ogrisel mentioned this pull request May 12, 2026

FIX: fix array API support for float32-only devices on integer inputs scikit-learn/scikit-learn#33898

Merged

ev-br force-pushed the f32_device branch from b6334c8 to 8d4eb67 Compare May 25, 2026 10:22

ev-br changed the title ~~WIP: add a float32-only device~~ ENH: add a float32-only device May 29, 2026

ev-br mentioned this pull request Jun 23, 2026

Add DLPack device number to arrays #212

Merged

ev-br and others added 4 commits June 30, 2026 17:26

add an f32-only device, move device handling to _devices.py

2c489ca

refactor _info.py for device-specific dtypes

f321cdc

adapt creation functions for device-specific dtypes; add tests

94786bf

manually sync dlpack enum updates to _devices.py

0832120

As a part of a rebase/conflict resolution, this commit simply moves the dlpack/device updates to _devices.py. Co-authored-by: Tim Head <betatim@gmail.com>

ev-br force-pushed the f32_device branch from e17e42d to 3429bd4 Compare June 30, 2026 15:37

MAINT: fix up after the rebase

346c59a

ev-br force-pushed the f32_device branch from 3429bd4 to 346c59a Compare June 30, 2026 15:44

add dlpack_device dunder to the F32_device

dbe38b8

betatim reviewed Jul 6, 2026

View reviewed changes

ev-br and others added 3 commits July 6, 2026 16:45

Update array_api_strict/_creation_functions.py

826f8f2

Co-authored-by: Tim Head <betatim@gmail.com>

address review comments

df1574d

raise if (device_type, device_id) is not found in the DLPack mappings

a15c7bc

	_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "F32_device")
	_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "no_float64_device")

-            # find out the default dtype for the device
+            # The dtype selected by Numpy might not be the default dtype
+            # on this device. If the dtype is not supported by the device we
+            # try to find one that is, aka the default dtype of this device.

		return full_like(a, fill_value=42.0, args, *kwds)


		class TestDefaultDType:

Uh oh!

Conversation

ev-br commented Apr 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

betatim commented Apr 27, 2026

Uh oh!

ogrisel commented Apr 30, 2026

Uh oh!

ogrisel commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ev-br commented May 25, 2026

Uh oh!

ev-br commented Jun 30, 2026

Uh oh!

rgommers commented Jul 6, 2026

Uh oh!

betatim left a comment

Choose a reason for hiding this comment

Uh oh!

betatim Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

ev-br Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

betatim Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

ev-br Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

betatim Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

ev-br Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

ev-br commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ev-br commented Apr 24, 2026 •

edited

Loading

ogrisel commented Apr 30, 2026 •

edited

Loading