[SPMD] Allow dumping post optimizations hlo (#5302)
Summary:
This pull request partial reverts the change in #5266 to re-enble
dumping post optimizations hlo.
Test Plan:
XLA_USE_SPMD=1 PJRT_DEVICE=TPU python test/spmd/test_xla_sharding.py -v -k test_xla_sharded_hlo_dump_post_optimizations
