[Bug] Fix/sglang jax support pathways

[aolemila]

Resolves #861

It's a good idea to open an issue first for discussion.

Add pathwaysutils.initialize() and use utils.show_hbm_usage = utils.show_hbm_usage.

Reference

Colab Notebook

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

Results

Environment

# Pathways
JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

[device(18,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(19,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(20,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(21,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3)]

# run 
JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python3 scripts/grpo_demo_sglang_jax_rollout.py