Enable JAX memory tests for GPUs and CPUs
PjRt GPU and CPU has recently gotten memory space support with just one memory space per device, so enabling relevant JAX memory tests. Most tests cannot be enabled yet because they rely on `unpinned_host`, so only enabling `ShardingMemoriesTest` for now.
PiperOrigin-RevId: 633335638
