Rob Suderman

Address comment linked below. `jit_kernel` is more an implementation detail and should not be exposed to the user. See about excluding it from import / reworking to be non-visible. https://github.com/iree-org/iree-jax/blob/main/iree/jax/builtins.py#L46

Serialize Executables crashing when compiling LLaMa on async-cpu

21

The following dispatches appear to cause a crash when compiling a llama model. Unrolling / vectorization makes 20K+ lines of generated code which likely causes final LLVM compilation to completely...

Expand `arith.minf` and `arith.maxf` expand ops for non-`f32` types

9

x86 compilation fails with `minf` and `maxf` reduction operations. This is the ir for the corresponding failing test: ``` func.func @main(%arg0 : tensor, %arg1 : tensor) -> tensor { %0...

Convert `*.rng_bit_set_generator` to use `i32` state bits before iree-util-demote-i64-to-i32

1

We have had multiple failures due to iree-util-demote-i64-to-i32 being set by default. Before we execute this pass we should see about converting all cases of `*.rng_bit_generator` to use an `i32`...

[Epic] Add support for `mhlo.custom_call`

4

There are some operations that we will likely support via library calls (e.g. QR, Eigh, etc). We should add support for a method to call out to related libraries to...

Determine how to support JAX's callback system

3

JAX supports registering callbacks that can be triggered via the backend environment. This may require implementation via the `mhlo.custom_call` system or an alternative integration method via the PJRT events system.

These operations can directly be bitcasted between types for iree.

We can support `onnx.Adam` by decomposing into elementwise operations. The numerics slightly differ from `onnxruntime` however they appear to be within margins.

This restriction is a bit needless and could be wrong for larger memory blocks.

[onnx] Fix `onnx.Slice` for dynamic axis

1

If the axis value is dynamic we need to iterate through the options and match to the correct dim.