More fixes for 2023.12 support #166

asmeurer · 2024-07-29T19:07:20Z

Fixes #152

- Ensure the arrays that are created are created on the same device as x. (fixes data-apis#177) - Make clip() work with dask.array. The workaround avoid uint64 -> float64 promotion does not work here. (fixes data-apis#176) - Fix loss of precision when clipping a float64 tensor with torch due to the scalar being converted to a float32 tensor.

I'm not sure if all the details here are correct. See data-apis#127 (comment).

Some of these things have to be inspected manually, and I'm not completely certain everything here is correct.

asmeurer · 2024-09-11T21:51:46Z

array_api_compat/torch/_info.py

I'm hoping someone with more knowledge of pytorch can review this. I've made various assumptions here, which I'll try to outline in comments below. If anyone can confirm whether those assumptions are valid, that would be helpful.

asmeurer · 2024-09-11T21:52:15Z

array_api_compat/torch/_info.py

+        """
+        return {
+            "boolean indexing": True,
+            "data-dependent shapes": True,


I'm assuming boolean indexing and data-dependent shapes (i.e., functions like unique) always work in pytorch.

Could you clarify what you mean by always work? unique and nonzero will generally work in PyTorch the same way they do for NumPy, considering normal eager execution.

With the compiler, maybe depending on the opinions chosen, is a graph break considered working?

Complex dtypes will not work, but not due to issues with data dependent shapes, it is because our implementation will unconditionally sort the input.

I hadn't even considered the torch compiler. But I think as long as it functions without an error, that should be considered working (if that's not the case, then it might actually be worth flipping this flag in a compiler context, assuming that's easy).

These flags exist for libraries like dask or JAX that straight up don't allow these types of operations.

In that case I would say you are probably correct the way you have it. Those are intended to be supported and lack of support/conditional support in the compile context is considered a deficiency.

asmeurer · 2024-09-11T21:53:33Z

array_api_compat/torch/_info.py

+        'cpu'
+
+        """
+        return torch.device("cpu")


I'm assuming the default device in pytorch is always "cpu" (note there is currently an unfortunate distinction between the "default device" and "current device" in the standard. See data-apis/array-api#835). By "default device", I mean the device that is used by default when pytorch is first started.

Yes, this is correct.

asmeurer · 2024-09-11T21:54:17Z

array_api_compat/torch/_info.py

+         'indexing': torch.int64}
+
+        """
+        default_floating = torch.get_default_dtype()


If I'm understanding the docs correctly, this function will always give the default floating-point dtype (i.e., what is used by default for something like torch.ones() with no dtype argument).

Yes, this is true. However that "default" in torch is not static. It can be changed by the user, so that would be the "current default" and the "default default" would be torch.float32

Just a note to myself: I checked and torch does correctly fail if you set the default dtype to float64 and try to create a tensor on a device that doesn't support float64:

>>> torch.set_default_dtype(torch.float64) >>> torch.asarray(0., device='mps') Traceback (most recent call last): File "<stdin>", line 1, in <module> TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

(unlike certain other libraries that silently map float64 back to float32)

so that would be the "current default" and the "default default" would be torch.float32

I'm not sure which this should use then. I think it should be the "current default", but the meaning of "default" is ambiguous. I mentioned this at data-apis/array-api#835 (comment)

asmeurer · 2024-09-11T21:55:14Z

array_api_compat/torch/_info.py

+
+        """
+        default_floating = torch.get_default_dtype()
+        default_complex = torch.complex64 if default_floating == torch.float32 else torch.complex128


The docs explicitly state that the default complex dtype always matches the floating dtype. https://pytorch.org/docs/stable/generated/torch.set_default_dtype.html

asmeurer · 2024-09-11T21:55:31Z

array_api_compat/torch/_info.py

+        """
+        default_floating = torch.get_default_dtype()
+        default_complex = torch.complex64 if default_floating == torch.float32 else torch.complex128
+        default_integral = torch.asarray(0, device=device).dtype


Is there a way to access this that doesn't require creating a tensor?

You can also just hard code it, we don't have a default for integers internally. If this were to change it would break bc and would need to be updated everywhere an integer tensor is produced, manually.

You mean it's always int64? I wasn't sure if this would be different for certain devices.

We don't have a concept of the default integer type. It would always need to be provided or the type is deduced from the argument that is what happens in this case, python ints can be larger than int32 so we use int64.

I mean that I do not know of any mechanism for changing the behavior you see here. Our deduction rules for creation from a non array type can be seen here. If the device does not support int64 I would expect you to see an error like you do with MPS and float64 (which I did not know would happen, I expect they had to do some work to make sure that error gets raised.)

Note that that layer of code is going to run "above" the dispatch to device specific logic (as in closer to python), so there is no way for the device to influence how that type is deduced, the device will be involved when the underlying tensor object is constructed where it will only see that a dtype argument has been provided.

asmeurer · 2024-09-11T21:56:39Z

array_api_compat/torch/_info.py

+            "real floating": default_floating,
+            "complex floating": default_complex,
+            "integral": default_integral,
+            "indexing": default_integral,


I'm assuming the default indexing type is always the same as the default integer dtype (the default indexing type is a concept in the array API for functions that return indexing arrays. For example, nonzero should return an array with the default indexing type https://data-apis.org/array-api/latest/API_specification/generated/array_api.nonzero.html)

asmeurer · 2024-09-11T21:57:05Z

array_api_compat/torch/_info.py

+        uint8 = getattr(torch, "uint8", None)
+        uint16 = getattr(torch, "uint16", None)
+        uint32 = getattr(torch, "uint32", None)
+        uint64 = getattr(torch, "uint64", None)


Are these the only dtypes that can be undefined (I know newer torch versions have them, but I want to make sure older versions work here too).

Also I know the support for some of these is limited. Would it be more correct to always omit them, even when they are technically defined here? They aren't really fully supported from the point of view of the array API standard.

I would say that depends on the purpose of the declaration. If it is to list the definitions for data types provided by the library leave them. If the point is to declare the data types supported by the array api then drop them.

Yeah, I should probably remove them then. Too many array API things don't actually work with them, and this API is supposed to be a better way to check that than hasattr(xp, 'uint64').

asmeurer · 2024-09-11T21:58:06Z

array_api_compat/torch/_info.py

+                del res[k]
+                continue
+            try:
+                torch.empty((0,), dtype=v, device=device)


Is this the best way to test if a dtype is supported on a given device?

What defines supported?

If dtype and device are valid this should always work, storage is untyped so the dtype is only used to compute the number of bytes needed.

Any given operator may or may not correctly dispatch for the underlying dtype. If the type is not explicitly handed for a given operator it will throw.

asmeurer · 2024-09-11T21:58:44Z

array_api_compat/torch/_info.py

+            return res
+        raise ValueError(f"unsupported kind: {kind!r}")
+
+    @cache


I'm assuming it's safe to cache the output of this function (it's potentially expensive since it constructs tensors on the given device to test if a dtype is supported).

There is currently no way to define a new dtype than adding an entry to a enum in the source code.

asmeurer · 2024-09-11T21:59:05Z

array_api_compat/torch/_info.py

+                del res[k]
+        return res
+
+    @cache


I'm assuming the set of legal devices never changes at runtime and can be cached.

You can use hooks to register a custom device backend dynamically at runtime. I am unsure if this will simply add a new accepted device type string, or if the privateuseone string is overwritten by it.

asmeurer · 2024-09-11T22:00:44Z

array_api_compat/torch/_info.py

+        # currently supported devices. To do this, we first parse the error
+        # message of torch.device to get the list of all possible types of
+        # device:
+        try:


This is the big one. This is code is based on some discussions with @pearu. To get the list of possible torch devices, we first parse an error message, then check which of those devices actually work. If there is any better way of doing this, please let me know. I would definitely prefer if this functionality were built-in to pytorch (ditto for the other functions here too).

This is unfortunate. I think there might be a less expensive, but more ugly way to do this. Would adding something to surface this info even help, or would you still need to have something like this to support older versions?

Adding something would definitely be helpful. Note that ideally, pytorch will implement this exact API here, since it's part of the array API.

I noticed that pytorch sort of has APIs to do this better, e.g., https://pytorch.org/docs/stable/generated/torch.mps.device_count.html#torch.mps.device_count, but they are not consistent across all device types, and I didn't want to hard-code all the possible device types here since torch seems to support a lot of them.

Note: I think that API should exist for optional devices. And you can find the module programmatically with torch.get_device_module(<name>). The always available devices are of course special, but I can look into handling this on the pytorch side.

asmeurer · 2024-09-12T20:16:46Z

array_api_compat/cupy/_info.py

+         'int64': cupy.int64}
+
+        """
+        # TODO: Does this depend on device?


@leofang is it possible for cupy to not support some of the array API dtypes depending on a given device?

See the discussion at data-apis#166 (comment)

…ion)

asmeurer · 2024-09-30T21:35:38Z

So a couple of small things are still missing here #127. Most notable is repeat which is missing for PyTorch and has some minor issues with NumPy. There's also some concerns with default_device in the inspection API which might end up changing (see data-apis/array-api#835). But there are quite a few changes here so I want to get them merged and released. I will hold off on bumping the __array_api_version__ for now, though.

Remove floating-point promotion from sum, prod, and trace

0734064

Fixes data-apis#152

asmeurer mentioned this pull request Jul 29, 2024

Support for 2023.12 #127

Open

18 tasks

asmeurer added 11 commits July 29, 2024 13:09

Wrap trace for dask

25da0e4

Wrap hypot in torch

e383441

Add unstack to all wrapped libraries

5eb7abf

Run the test suite against the 2023 version of the standard on CI

e70bcc8

Update xfails for failing 2023 tests

d751db6

Merge branch 'main' into more-2023

b96e84b

Remove clip from dask xfails

2995a0f

Update numpy 1.21 xfails

284bd99

Add NumPy inspection namespace

11cb6ef

Add CuPy inspection APIs

4c9dd0e

I'm not sure if all the details here are correct. See data-apis#127 (comment).

ogrisel mentioned this pull request Sep 2, 2024

Automatically use the correct device in xp.clip with passed Python number literal as bounds #177

Closed

asmeurer added 3 commits September 11, 2024 15:01

Merge branch 'main' into more-2023

8e3f0b6

Add inspection namespace for torch

d3c4b3c

Some of these things have to be inspected manually, and I'm not completely certain everything here is correct.

Handle torch versions that do not have uint dtypes

be4fa68

asmeurer commented Sep 11, 2024

View reviewed changes

asmeurer mentioned this pull request Sep 12, 2024

Clarify definitions of "default device" and "current device" data-apis/array-api#835

Open

Remove uint16, uint32, and uint64 from the pytorch dtypes() output

774d175

asmeurer commented Sep 12, 2024

View reviewed changes

asmeurer added 12 commits September 12, 2024 14:26

Add inspection namespace for dask

231ef95

Add cumulative_sum wrapper for the numpy-likes

92ebdff

Remove a bunch of 2023.12 xfails

176a66a

Hard-code the default torch integral type to int64

46a2227

See the discussion at data-apis#166 (comment)

Ignore bare except in ruff checks

ab822f5

Add a comment

cb9acd4

Add cumulative_sum to torch

c0dd5b0

Fix the tests

1b47d96

Add repeat to the torch xfails

e472dcb

Update torch xfails

470e41a

Update numpy dev xfails

198d6d7

Add maximum and minimum torch wrappers (for fixed two-arg type promot…

ef7ad7a

…ion)

asmeurer enabled auto-merge September 30, 2024 21:35

Add dlpack xfails to numpy-dev

9d2e283

asmeurer merged commit b30a59e into data-apis:main Sep 30, 2024
43 checks passed

More fixes for 2023.12 support #166

More fixes for 2023.12 support #166

Conversation

asmeurer commented Jul 29, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

amjames Sep 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

asmeurer commented Sep 30, 2024

asmeurer commented Jul 29, 2024 •

edited

Loading

amjames Sep 12, 2024 •

edited

Loading