Enable entropy on multihost CPUs.

[RoshaniN]

• Standalone checkpointer was not working well on multihost CPUs. Recommendation was to refer to JAX arrays using with jax.spmd_mode('allow_all'):
• Simplified entropy addition.

Thanks for debugging with me @gobbleturk , thanks @raymondzouu for reporting this!

Tested using -

python3 xpk.py workload create --command "JAX_PLATFORMS=cpu python3 MaxText/standalone_checkpointer.py MaxText/configs/base.yml \
base_output_directory=gs:https://maxtext-experiments-multipod dataset_path=gs:https://max-datasets-rogue steps=120 checkpoint_period=30 enable_checkpointing=True async_checkpointing=False per_device_batch_size=1 ici_data_parallelism=1 ici_fsdp_parallelism=64 hardware=cpu;" \
--cluster test2-n2-standard-32-32 \
--docker-image=gcr.io/tpu-prod-env-multipod/roshanin_runner \
--num-slices=1 \
--device-type=n2-standard-32-64 \
--zone=us-central1-b \
--workload=roshanin-ps-12

[RoshaniN]

Seeing unrelated test failure https://screenshot.googleplex.com/zf7vN5462q6aU4T

[RoshaniN]

Seeing unrelated test failure https://screenshot.googleplex.com/zf7vN5462q6aU4T

#629 should fix the error.