apivovarov/hlo_convert_torch_xla.py

Created August 23, 2024 21:47

Star (1) You must be signed in to star a gist
Fork (0) You must be signed in to fork a gist

Select an option

Learn more about clone URLs
Clone this repository at <script src="https://gist.github.com/apivovarov/bbaf7d595449f3d2f3565e776d1a8d23.js"></script>
Save apivovarov/bbaf7d595449f3d2f3565e776d1a8d23 to your computer and use it in GitHub Desktop.

Download ZIP

Revisions

apivovarov created this gist Aug 23, 2024.

hlo_convert_torch_xla.py

-Original file line number
+Diff line change
@@ -0,0 +1,7 @@
+    import torch_xla.core.xla_builder as xb
+    fname = "model.hlo.pb"
+    with open(fname, mode="rb") as f:
+      comp = xb.computation_from_module_proto("foo", f.read())
+    print(xb.get_computation_hlo(comp))