Keith Rush issues

Repositories
Issues
Comments

Results 2 issues of


                                            Keith Rush

Corner-case sharding regression when replacing concrete mesh with abstract mesh

### Description With no data dependency on inputs, abstract mesh and concrete mesh in `jax.lax.with_sharding_constraint` result in different behavior, with concrete mesh behaving as expected. Sending a PR with this...

bug

vmap(SPMD axis)/shmap/(vmap with capture) pattern breaks batching

### Description By doing sufficiently-exciting capture + sharding' + mapping behavior, it is possible to induce jax's batching to witness inconsistent sizes for the batch axis. The following code snippet...

bug