Add option to dump hlo

[gobbleturk]

Description

• Add option to Dump HLO in MaxText
• This uploads soon after compilation for both train.py and train_compile, which is beneficial as compared to using --debug-dump-gcs on XPK which uploads only at the job end (which will never happen if the job hangs).
• This does have the drawback of assuming a GCS endpoint for the dump location, some code modification is required for arbitrary endpoints

Tests

Ran a few manual tests of train.py different dump options on both local devbox and on a larger XPK job, as well as some local train_compile tests.
We can modify our hybrid sim flow to use this so it will be covered in e2e test. It is hard to unit test because

1. it is not hermetic - depending on the state of the jax cache and gcs buckets
2. The XLA flags does not play nicely with pytest, I think only one backend will be set with an intial set of XLA flags, we cannot use different flags in different tests

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.