[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