Skip to content

WIP: add a float32-only device#206

Open
ev-br wants to merge 12 commits intodata-apis:mainfrom
ev-br:f32_device
Open

WIP: add a float32-only device#206
ev-br wants to merge 12 commits intodata-apis:mainfrom
ev-br:f32_device

Conversation

@ev-br
Copy link
Copy Markdown
Member

@ev-br ev-br commented Apr 24, 2026

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make one of the devices have a lower maximum float precision Allow changing the default dtypes

1 participant