Find a way to get sharding information to jax from StableHLO code through a protobuf

[ajakovljevicTT]

Currently, we pass dummy sharding (replicated on all devices) to our inputs, as jax needs to be sharding aware when executing a graph. The proper way to do this is to get this information from the StableHLO compile through protobuf or xla::CompileOptions, which would require linking the whole xla repo. There needs to be more exploration on this topic.