Adding implementation for PJRT_Executable_OutputElementTypes and PJRT_Executable_OutputDimensions

[sdjukicTT]

Ticket

#164

Problem description

PJRT_Executable_OutputElementTypes and PJRT_Executable_OutputDimensions were unimplemented.

What's changed

In LoadedExecutableInstance::Execute we got the shape for every PJRT_Buffer output from binary.getProgramOutputs and image_->isOutputScalar. Now we store the shape in ExecutableImage class and have a getter for the final shape, and not isOutputScalar.

Adding a vector that stores PJRT_Buffer_Type for every output in ExecutableImage, as well as a getter, for the implementation of PJRT_Executable_OutputElementTypes.

PJRT_Executable_OutputDimensions requires the dimensions to be in a different format. It returns an array of output dimension sizes and an array of concatenated dimensions, so we also store that type of pointer in ExecutableImage. We store it because we are responsible for memory deallocation, and not the caller of PJRT_Executable_OutputDimensions. But because concatenated dimensions are rarely needed, these fields are nullptr until get_output_dims_concatenated is called.

Notes

It is possible that PJRT_Executable_OutputElementTypes should return only different unique PJRT_Buffer_Types, and not a type for every output.

We might want to store information about stride of every output in ExecutableImage as well, so binary.getProgramOutput wouldn't have to be called in LoadedExecutableInstance::Execute, and all the information about outputs would be at one place.

Checklist

• New/Existing tests provide coverage for changes

[github-actions]

	Tests	Passed ✅	Skipped ⚠️	Failed
TT-XLA Tests	599 ran	431 passed	168 skipped	0 failed

Test	Result
No test annotations available

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 36.11111% with 23 lines in your changes missing coverage. Please review.

Project coverage is 76.92%. Comparing base (abe6da2) to head (fa67c0a).
Report is 4 commits behind head on main.

✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/common/pjrt_implementation/executable_image.cc	32.35%	23 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #238      +/-   ##
==========================================
- Coverage   77.30%   76.92%   -0.39%     
==========================================
  Files          21       21              
  Lines        1018     1053      +35     
==========================================
+ Hits          787      810      +23     
- Misses        231      243      +12     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ajakovljevicTT]

Describe this in a bit more detail, something like: "Check whether the information on output dimension has been populated".