Skip to content

Conversation

@matthiasdiener
Copy link
Collaborator

@matthiasdiener matthiasdiener commented May 23, 2025

Needs inducer/arraycontext#315 (perhaps)

Please squash

@matthiasdiener matthiasdiener self-assigned this May 23, 2025
@matthiasdiener matthiasdiener marked this pull request as draft May 23, 2025 16:42
@matthiasdiener matthiasdiener marked this pull request as ready for review May 28, 2025 22:34
@matthiasdiener matthiasdiener requested a review from inducer May 28, 2025 22:34
@matthiasdiener matthiasdiener changed the title add MPIEagerJaxArrayContext add MPIEagerJAXArrayContext May 28, 2025
Copy link
Collaborator

@alexfikl alexfikl left a 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 😁

Comment on lines 571 to 572
grp_field = signed_face_ones_numpy[igrp]
sign_mask = np.ones_like(grp_field)
Copy link
Collaborator

@alexfikl alexfikl May 29, 2025

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_ones and signed_face_ones_numpy things.
  • 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?

Copy link
Collaborator Author

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).

Copy link
Collaborator

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?)

@matthiasdiener
Copy link
Collaborator Author

Closing in favor of #388.

@matthiasdiener matthiasdiener deleted the mpieagerjax branch September 10, 2025 19:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants