Check and update return sharding

[wooseokTT]

Ticket

#2297

Problem description

PJRT expects the return values to be sharded instead of being merged into a tensor with a single buffer if there is sharding attribute in return value.

What's changed

This PR provides special care for the return values just like we do for input tensors. It checks the return value attribute and if the sharding attribute is there, mark the returning mesh_shard op as manual such that runtime will not perform any concat operation and returns multiple buffers.

Checklist

• [V] New/Existing tests provide coverage for changes

[nsmithtt]

I think just make sure to fix the ifdef because that can lead to confusing compiler errors. Lmk what you think about the enum naming otherwise lgtm! Thanks