jax
0b5d885e - Handle with_sharding_constraint inside arbitrary subjaxprs. (#3339)

Commit

6 years ago

Handle with_sharding_constraint inside arbitrary subjaxprs. (#3339) Also fixes a bug where a replicated sharding constraint would incorrectly trigger an error.

References

#3339 - Handle with_sharding_constraint inside arbitrary subjaxprs.

Author

skye

Parents

29740de6

jax 0b5d885e - Handle with_sharding_constraint inside arbitrary subjaxprs. (#3339)

jax
0b5d885e - Handle with_sharding_constraint inside arbitrary subjaxprs. (#3339)