Python Array API Compatibility Tracker

[pmeier]

Progress tracker for all compatibility issues with Python Array API specification:

• __array_namespace__ (add at/towards the end, it's the attribute that declares compliance)
• Support for uint16, uint32, and uint64 #58734
• type promotion with 0d-tensors diverges from array API specification #58736
• Python Array API: support constants #58739
• support for scalars in logical_(and|not|or|xor) #58740
• torch.size() diverges from array API specification #58741
• Python Array API New Operators Tracker #58742
• Implement improved DLPack support #55090
• torch.(min|max)(..., dim=...) diverges from array API specification #58745
• torch.nonzero diverges from the specification (it returns a tensor, not a tuple, by default) - see [numpy] torch.nonzero is similar to np.argwhere not np.nonzero #64502
• Support negative step sizes for slicing #59786
• Remove support for multiple ellipses in slicing #59787
• Bitwise shifts should retain the same dtype as the first input #59867
• Python scalars should be promoted to the same dtype as the respective tensor #59868
• torch.asarray does not detect dtype of Python scalars #70591
• full should take an integer size #70906
• finfo(...).tiny should be aliased to finfo(...).smallest_normal #70909
• eye should support other diagonals than the main one #70910
• support setting step in arange without setting end #70914
• arange should return empty array if bounds are inconsistent with step sign #70915
• uint8 scalar tensors cannot be used for integer indexing #70916
• make torch.(ceil|floor|round|trunc) no-ops for integer inputs #70918
• linspace should support an endpoint parameter #70919
• unique should be split into four partial functions #70920
• sort should only return the sorted input #70921
• argsort is missing the stable parameter #70922
• torch.where to support Python scalars and type promotion #9190
• change supported arguments for parameter dim in squeeze #70924
• allow dim=None in concat #70925

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411

[leofang]

cc: @kmaehashi @emcastillo @asi1024

[pmeier]

We've investigated each API in the Data API specification and summarized modifications needed to run code written using the API on PyTorch.
Here is the full list.
https://docs.google.com/spreadsheets/d/13amsIeHhe38apCeZrOmsbsHT-aysmYur9-cugw5jVg0/edit#gid=0
In the list, we have classified each API into 6 categories:

• No Changes: The current PyTorch implementation already conforms to the Data API.
• New API: Need to add a new function or constant.
• Enhancement: Need enhancement to the existing function (add new argument, etc.).
• Enhancement with BC: Need to change signature of the existing function. (breaking change if we implement the API directly under torch.* namespace)
• Behavior: Need to change behavior of the existing function. (breaking change if we implement the API directly under torch.* namespace)
• N/A: The detailed specification is not determined yet in the Data API.

To implement the Data API directly under torch.* namespace, we need to discuss how to deal with the breaking changes. Major topics are:

• Arguments: Some API cannot be compatible with the Data API specification without breaking changes. e.g., torch.linspace(start, end, steps, ...) v.s. Data API linspace(start, stop, num, ...).
• Argument names: PyTorch APIs conventionally use dim/dims as an argument name (e.g., torch.stack), whereas the Data API requires axis keyword argument.
• Return value: For example, torch.nonzero returns a tensor when as_tuple=False (default) whereas the Data API requires to always return as tuple.
• Behavior:
  • torch.cross looks for the first dimension found with the size 3, whereas the Data API uses -1 as default axis.
  • floordiv operator needs be changed to take floor. (The documentation says This behavior will be deprecated in a future PyTorch release so it may be a good time to change the behavior)

Originally posted by @asi1024 in #54581 (comment)

[asi1024]

@pmeier Sounds good to me! Thank you!

[rgommers]

I have added a module: python array api label to gh-24422, so the bot will Cc people listed when that label gets added to new issues. Easier than manually Cc-ing a bunch of people. If you want to get added and don't have permissions to add yourself via editing that issue description, please let me know and I'll add you.

[asmeurer]

@rgommers can you add me?

[leofang]

+1