fix util dump_hlo_graph with new jax API

pfackeldey · Oct 29, 2024 · da4c737 · da4c737
1 parent 2d31648
commit da4c737
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/evermore/util.py b/src/evermore/util.py
@@ -167,4 +167,4 @@ def f(x: jax.Array) -> jax.Array:
         filepath = pathlib.Path('graph.gv')
         filepath.write_text(dump_hlo_graph(f, x), encoding='ascii')
     """
-    return jax.xla_computation(fun)(*args, **kwargs).as_hlo_dot_graph()
+    return jax.jit(fun).lower(*args, **kwargs).compiler_ir("hlo").as_hlo_dot_graph()  # type: ignore[union-attr]