[Shmap/PartialAuto] Temporary solution for debug.print inside a partial-auto shard map.

[yliu120]

This commit needs the XLA commit: openxla/xla#20944 to work.

However this is not the ultimate fix for this. We want some solution that works in very short term as this feature is crucial for debugging.

The issue with debug.print not working inside a partial-auto shard map is because,

1. The tensor will be further partitioned in GSPMD however the host callback is created during IR lowering time with the shape traced in shard map.
2. based on 1), a permanent fix would be serializing the args info into the opaque string and after SPMD partitioning, XLA should pack the partitioned args info into the backend config opaque string.