@@ -0,0 +1,986 @@
+# RUN: %25PYTHON %25s%0A# XFAIL: *%0A%0Aimport iree.compiler%0Aimport iree.jax%0Aimport iree.runtime%0Aimport jax%0Aimport jax.numpy as jnp%0A%0A# https://github.com/iree-org/iree/issues/10230%0A%0Aarg0 = jnp.array(%5B%5B3.0, 2.0, 4.0%5D, %5B5, 6, 7%5D, %5B3, 4, 5%5D%5D)%0Aarg1 = jnp.array(%5B%5B3, 2, 4%5D, %5B5, 6, 7%5D, %5B3, 4, 5%5D%5D)%0A%0Aclass BugTest(iree.jax.Program):%0A  @iree.jax.kernel%0A  def _main(x, y):%0A    t = jnp.einsum(%22ki,jk,im,in,ix-%3Eijmnx%22, x, jnp.einsum(%22ik,kj-%3Eij%22, x, y) + jnp.power(x + y,3), y,x,y)%0A    return jax.nn.softmax(t)%0A%0A  def main(self, x=iree.jax.like(arg0), y=iree.jax.like(arg1)):%0A    return self._main(x, y)%0A%0Air = str(iree.jax.Program.get_mlir_module(BugTest()))%0Aprint(ir)%0A%0Airee_config = iree.runtime.system_api.Config(%22local-task%22)%0Airee_binary = iree.compiler.compile_str(%0A    ir, target_backends=%5B%22llvm-cpu%22%5D, input_type=%22mhlo%22)%0Avm_module = iree.runtime.VmModule.from_flatbuffer(instance, iree_binary)%0Amodule_object = iree.runtime.load_vm_module(vm_module, iree_config)%0Aout = module_object%5B%22main%22%5D(arg0, arg1)%0A

Add bug test for 10230 (#79)