-
Notifications
You must be signed in to change notification settings - Fork 18
add MPIEagerJAXArrayContext #380
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
Conversation
714d73c to
0ea000f
Compare
0ea000f to
60fbd73
Compare
alexfikl
left a comment
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.
Left a few comments for that _signed_face_ones function 😁
grudge/geometry/metrics.py
Outdated
| grp_field = signed_face_ones_numpy[igrp] | ||
| sign_mask = np.ones_like(grp_field) |
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.
I have a hard time remembering why this was written this way (I think those signed_face_ones arrays are just created this way so that they have the right shape?), but maybe we can improve it a bit.
How about this?
- Remove the
signed_face_onesandsigned_face_ones_numpythings. - Get the discretization
discr = dcoll.discr_from_dd(dd.with_discr_tag(DISCR_TAG_BASE)) - For each zip(discr.groups, conn.groups),
mask = np.ones((dgrp.nelements, dgrp.nunit_dofs), dtype=discr.real_dtype) - The rest just stays the same (?)
Would something like that work? The main idea being that we'll just create everything in numpy and then transfer it over to the array context, so we don't have to worry about how writable the arrays are.
We might want to tag the end result in the same way that Discretization.zeros does too?
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.
Thanks for the suggestion! What do you think of 85a0d54? (feel free to push directly to this branch if I did something silly, I am very unfamiliar with this code).
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.
Looks good to me! (doesn't look like the test failure is related?)
Co-authored-by: Alex Fikl <[email protected]>
|
Closing in favor of #388. |
Needs inducer/arraycontext#315 (perhaps)Please squash