Skip to content

Instantly share code, notes, and snippets.

@apivovarov

apivovarov/hlo_convert_jax.py

Created August 23, 2024 21:46
Show Gist options
  • Save apivovarov/be865089bfa73b2134cee24e217cf27f to your computer and use it in GitHub Desktop.
Save apivovarov/be865089bfa73b2134cee24e217cf27f to your computer and use it in GitHub Desktop.
Download ZIP
[JAX] convert hlo.pb to text hlo
Raw
hlo_convert_jax.py
from jax.lib.xla_bridge import xla_client
fname = "model.hlo.pb"
with open(fname, mode="rb") as f:
comp = xla_client.XlaComputation(f.read())
print(comp.as_hlo_text())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment