search

data-apis / array-api-compat

Compatibility layer for common array libraries to support the Array API
https://data-apis.org/array-api-compat/
MIT License
76 stars 26 forks source link

PyTorch compatibility layer #14

Closed asmeurer closed 1 year ago

asmeurer commented 1 year ago

This looks quite good. I left a number of comments, mostly with suggestions for improvements. In general, it'd be good to add tests (and linting!) for the issues that this review may have found as well.

Yes, I want to add the test suite to the CI. I will finish #13 before releasing this. I guess we can add some basic linting, although we'll need some ignore comments due to the import * stuff that I have to do in the init files.

Now, why are there some functions that preprocess their inputs with torch.asarray and others that don't?

I did that in the functions that I needed to access attributes on the input (x.ndim). Otherwise I just pass everything through to the corresponding torch functions. It's quite possible I missed some cases here as the test suite is only going to be testing inputs that are valid according to the spec (i.e., only passing array inputs to functions).

A comment on reductions. All the reductions that don't support multiple axes may be implemented via a call to torch.movedim, moving all the relevant dimensions to the end, then a call to flatten and a call to the reduction with axis=-1. This would incur in one copy and one call to the kernel, so it should be much more efficient in eager mode, and much easier to optimise for the compiler. It should be easy to do this in a generic way via a wrapper.

I'll take a look.

asmeurer commented 1 year ago

This is ready for final review. I have addressed all the review comments. I have also added CI with XFAILs/SKIPs (some tests have to be skipped because they crash on CI), and have updated the README with a list of the known issues. Assuming everything looks good here, I can cut a release with this.

lezcano commented 1 year ago

how come this one was merged before the final review?

asmeurer commented 1 year ago

Because I am on PTO this week and I wanted to get this merged and released before then. If you have any further comments feel free to comment here or open a new issue.