[PJRT:GPU] Fail on unknown memory space for execution output buffers

[jaro-sevcik]

This PR addresses @yliu120 's post-landing comments from #14088:

1. Error out on unknown memory_space in the output's layout (this requires a bit of plumbing to propagate the error).
2. Use memory_space_by_kind_id in place of memory_space_by_kind (this requires some ugly cast, not sure if it's worth it - performance should not matter much here).
3. Avoid shape copy, pass by reference.

[dimitar-asenov]

This changes here are fine, but there are no test changes. While the changes are small, they do include API changes, and I'm a bit surprised that no test needs to change. Could you please add a test?

[jaro-sevcik]

This changes here are fine, but there are no test changes. While the changes are small, they do include API changes, and I'm a bit surprised that no test needs to change. Could you please add a test?

This error path is not easy to test because you would need to create an executable with output memory space that are not present on device. I think this might only be possible by hacking serialized executable. Is that what you have in mind?

As for the API change, this is a private method that is not used anywhere else (but perhaps it is used/overridden in PjRtTpuClient or InternalPjRtTpuClient?).

@yliu120 , could you please clarify what you had in mind in #14088 (comment)?

[dimitar-asenov]

OK, I see that the API is actually really an internal thing, and this change covers all the uses.