jax

266e6398 - [Mosaic TPU] Add support for slicing in reshape->store fusion optimization