JAX to MLIR using IREE example

jax_to_mlir.py

import jax
import iree.compiler as compiler

def YOUR_FUNCTION(x):
    return x + 1

input_sample = [
    jax.numpy.array([0.,])
]

aot = jax.jit(YOUR_FUNCTION).lower(*input_sample)
hlo_proto = aot.as_text()

mlir = compiler.compile_str(
        hlo_proto,
        target_backends=['llvm-cpu'],
        output_generic_mlir=True,
        input_type=compiler.InputType.STABLEHLO_XLA,
        output_format=compiler.OutputFormat.MLIR_TEXT,
        strip_debug_ops=True,
        strip_source_map=True,
        output_mlir_debuginfo=False,
        )

print(mlir.decode('utf8'))