jax
fa737e58 - [Pallas] Move CoreMemorySpace to Pallas core

Commit
15 days ago
[Pallas] Move CoreMemorySpace to Pallas core Why? mpmd_map is not TPU specific and we will need to extract CoreMemorySpaces from aval inputs to extend the axis env. We can't do this if CMS is part of pltpu. There's no inherent reason it needs to be. PiperOrigin-RevId: 903410710
Author
Parents
Loading