-
Notifications
You must be signed in to change notification settings - Fork 60
Array API support for core functions #844
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
2e5e20c
to
6ce7346
Compare
I just implemented the sphere. I implemented the inside/outside with masks. While, as noted in the comments, this might not be ideal for CPU(I'm not quite sure about this, but I assume you investigated that and found it to be true), It will make a good difference for the gpu. Ideally, these core functions should be implemented fully with array operations, so that backends like |
6ce7346
to
6db8100
Compare
Also note that each commit is individually reviewable, and I plan to keep it that way :) |
Additionally, if someone is willing to work on #732, I would feel a lot more confident in my translations here. |
a2d4c72
to
5a3e6f2
Compare
Added cuboid, polyline, and triangle. |
b3cfa8a
to
f928a32
Compare
Hi @purepani, Thank you very much for starting this ! Not sure yet if we should split this in multiple PRs but in the first step, modifications should only apply to the core functions ! Later when SciPy functions (Rotation) is implemented in an array-api compatible way we can start working the the general functional interface. And finally see if this can be applied all the way through the object oriented interface. BR, Alex- |
I have the commits organized so that it's super easy to create multiple PRs if needed(you can see the merge commit I have; they are already separate branches), so I can make them super easily if you'd prefer that! Or I can leave it all here! Just let me know if you want one or the other. Also, right now, they change the core functions api out of some necessity. I will have to think about how to improve that. |
Hey @purepani, Concerning the FE checks #732: Not really relevant. Our implementations are all properly tested - the FE tests were simply thought as a "second guarantee" that things are ok. In any case FE is massively lacking the precision the analytical method offers. I think @Alexboiboi is right that we should only start on the core functions for now, top level Concerning special functions: in some packages some are well implemented and ready in compiled form - in which case we should make use of them. In other cases we should have our own (which will never be as fast). Not 100% sure how to handle this ? br michael |
Just to respond to this, the current way the array api works is that any array Some of the things that we should look to the current efforts in scipy are listed here: https://data-apis.org/array-api/latest/design_topics/index.html |
70a1dcb
to
51b6193
Compare
... observers=xp.array([(1,1,1), (2,2,2)]), | ||
... dimensions=xp.array([(1,1,1), (1,2,3)]), | ||
... polarizations=xp.array([(0,0,1), (.5,.5,0)]), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
xp.array
does not exist in the standard, you probably want xp.asarray
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Same here, I thought I corrected them
... polarizations=xp.array([(0,0,1), (.5,.5,0)]), | ||
... ) | ||
>>> with np.printoptions(precision=3): | ||
>>> with xp.printoptions(precision=3): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this is not in the standard
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Huh, I thought I removed these
r = xp.sqrt(xp.vecdot(observers, observers))[ | ||
:, xp.newaxis | ||
] # faster than np.linalg.norm |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@Alexboiboi @OrtnerMichael @purepani of course feel free to proceed however you would like! But I would suggest the following integration strategy:
|
Yes, in a catchphrase the pattern this work enables is "array type in === array type out". I highly recommend reading scipy/scipy#18286 to understand the design we are implementing in SciPy. |
51b6193
to
39a63f7
Compare
I disabled the "Allow edits by maintainers" for now since I keep forgetting to run the precommit hook manually at the moment(I use jj which doesn't support git precommit hooks), and the automatic commit from the bot is a bit annoying. Obviously I'll try and format stuff when I remember to run ruff, but it's a bit annoying to deal with the bot adding commits every time I push. |
f4c5131
to
0575e2c
Compare
For getting elliptic functions working, I used The reason I chose to do this right now instead of writing a python implementation for supporting gpus was because The solution in jax would be to use Presumably that utility is something that would fit in |
hi @purepani Sorry for the missing feedback from our side. So far I was not able to review all your PR's. I'm about to complete a publication and will then have more time to invest...probably starting next week. As this touches the computational core I'll give it a very very thorough review. We should go through this function by function. which core function / PR do you suggest we start with ? |
Hey @OrtnerMichael After that, I've been throwing all my code in #861. This is just to get an overview of the changes, as I started separating out the code into more PRs. Most of the individual functions are ready for review, but I need to split out the code in the last commit. I'll probably have that ready once the core tests PR is merged. I'll update here, as well as mark the other PRs as ready once I do so. |
@OrtnerMichael I wrote a suggested review strategy here: #844 (comment) |
Hey @lucascolley |
@lithomas1 do you know if there is an issue or some other status somewhere for |
In an ideal world, yes. Unfortunately that guidance can't really hold for pre-existing array libraries. I guess we could look at some alternative for array-api-compat, but not sure what that would look like. |
Maybe a function in |
There is https://github.com/data-apis/array-api-compat/blob/main/array_api_compat/dask/array/linalg.py, but IIUC it is a WIP and somewhat limited by what is implemented in Dask proper. |
I realize my phrasing was slightly unclear. ...
if hasattr(xp, 'linalg'):
return xp.linalg.cross(x, y) we can do from array_api_compat import check_extensions, Extensions
if check_extensions(xp, Extensions.LINALG):
return xp.linalg.cross(x, y) (only using an enum in my example because I like type safety, but With it part of the spec if xp.has_extension('linalg'):
return xp.linalg.cross(x, y) or maybe if 'linalg' in xp.extensions:
print(xp.extensions) # ['linalg', 'fft']
return xp.linalg.cross(x, y) |
Ah, got you. Feel free to open an issue on the array-api repo! I guess until now, we have made do with external knowledge of the limitations of pre-existing libraries. Real examples of hitting this for missing features which are not likely to be added soon sound like good motivation for this sort of addition, however. |
5b5a533
to
5a2bc13
Compare
Ok @OrtnerMichael, I have some stuff that should be reviewable. I'll run the formatter later since my local branches are a bit of a mess and it's a bit tricky to resolve everything. So first, this PR just adds core tests, as mentioned earlier. This should be reviewed first: Then, here is a list of complete PRs, which work correctly with the all the backends that I added to the tests(dask, numpy, array_api_strict, and jax). I added tests for multiple backends since the lazy backends like jax are a bit more restrictive in what you can do with them.
Note that this is all the core functions except for the cylinder segment, since I think that core function will need to be refactored before a simple array api conversion can be done(I believe it's best to vectorize the application of all the different cases. Additionally, this PR contains all the changes to the core functions, and some in-progress changes: #861 Feel free to review the core functions in any order. The dipole is obviously the simplest, so that will probably be best to start off with. I think I'll need to split off the modifications of the tests to support the array api before merging, but otherwise, they should be suitably reviewable. |
I don't think so. There's not really anything actionable on the array-api-compat side (and there doesn't seem to be anyone working on implementing linalg stuff in dask for a while now) |
This PR will track the status of array api support per #792.
TODO:
np.
withxp.
andlen(r)
withr.shape[0]
, and maybe other things that I forgot about finding and replacing.getB
/getH
/getM
/getJ
xp.linalg
functions to be guarded under an if statement, or using a utility. Currently, the arrays are assumed to have the linalg extension, but this is not guaranteed by the api.xp.float64
for everything.lazy_while
for the elliptic integrals.So far, I've implemented support for the dipole core function, and also modified the test to use
array-api-strict
instead ofnumpy
.I had to change the core function to use masks to avoid division by 0 instead of ignoring warnings, since otherwise there is no way to ignore the warnings in a general way. I haven't tested this yet with any API backends.
I'm planning on using this PR to collect all the commits for the core functions, but I am also able to split them into separate PR, since that will be easier to review(hence the empty dev commit). As of right now, the dipole could be merged on it's on after review, since it doesn't change the behavior of the library(though an argument is added to the dipole function for the array backend). I can create a separate PR for each function as they are implemented if wanted, but I'll keep this draft PR to track it's overall progress.