jax

2fe54d1c - [pallas:mosaic] Added `CoreMemorySpace`

[pallas:mosaic] Added `CoreMemorySpace` This is a specialization of a `pltpu.MemorySpace` for a specific core type. For example, `pltpu.SMEM @ pltpu.CoreType.TC`. Not all memory space/core type combinations are valid. The lowering does not currently support `CoreMemorySpace`. This will be done in a follow up. PiperOrigin-RevId: 886153324