jax.xla_computation as_hlo_text()

[aoshujiaocheng]

import jax

def f(x): return jax.numpy.sin(jax.numpy.cos(x))
c = jax.xla_computation(f)(3.)
print(c.as_hlo_text())

How to convert c.as_hlo_text() into c?

[justinjfu]

JAX doesn't have python bindings for parsing HLO strings, however the machinery exists in XLA/C++. If you just want to run the function, you could perhaps try running the module using XLA's infrastructure (run_hlo_module - https://openxla.org/xla/tools).

Also as a heads up - jax.xla_computation will be deprecated with the upcoming 0.4.30 release.

[yashk2810]

Note 0.4.30 is released and jax.xla_computation is deprecated. Please use jax.jit(f).lower(3.).compiler_ir('hlo') as the replacement.