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
Add in_parts
and out_parts
optional arguments jax.xla_computation
.
#3771
Conversation
This allows partitioned computations in `xla_computation`, like those produced by `sharded_jit`.
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.
Seems reasonable! Maybe we could avoid the argument proliferation here in the future by having (a successor to) xla_computation
wrap the jitted, pmapped, or sharded_jitted function rather than standing in for the jit/pmap/sharded_jit.
Agreed, I like that idea. This is easier for now though. |
…mputation`. (google#3771)" This reverts commit dbc3f83. This is breaking some google-internal users of xla_computation. Reverting while I investigate.
…n`. (google#3771) This allows partitioned computations in `xla_computation`, like those produced by `sharded_jit`.
…mputation`. (google#3771)" (google#3780) This reverts commit dbc3f83. This is breaking some google-internal users of xla_computation. Reverting while I investigate.
…n`. (google#3771) This allows partitioned computations in `xla_computation`, like those produced by `sharded_jit`.
…n`. (google#3771) This allows partitioned computations in `xla_computation`, like those produced by `sharded_jit`.
…mputation`. (google#3771)" (google#3780) This reverts commit dbc3f83. This is breaking some google-internal users of xla_computation. Reverting while I investigate.
…n`. (google#3771) This allows partitioned computations in `xla_computation`, like those produced by `sharded_jit`.
…mputation`. (google#3771)" (google#3780) This reverts commit dbc3f83. This is breaking some google-internal users of xla_computation. Reverting while I investigate.
PAIR=skye PR google#3771 (reverted in google#3780) redux
PR google#3771 redux (reverted in google#3780) Co-authored-by: Skye Wanderman-Milne <skyewm@google.com>
This allows partitioned computations in
xla_computation
, like those produced bysharded_jit
.