Matthew Johnson

Before, on CPU: After, on CPU: Following up on #11866. I don't know if these small compile time differences matter, but hey, roofshots! And the jaxpr pretty-printing win is enough...

**Before:** **After:** No more `lt`, `add`, `select`, or `convert_element_type`! This might actually improve compilation time a bit; on CPU, **the tiny benchmark added went from ~15.5ms to 13ms on my...

add docstrings

2

Docstrings for some key functions and data structures, especially `Node` and `Box`, might help. Any suggestions for specific places to add docstrings?

Co-authored-by: Yash Katariya

fixes #10828

make config.jax_platforms = 'cpu' suppress no-cpu warning

5

also: * switch the warning to use `warnings.warn` rather than `absl.logging.warning` * fix a test which wasn't being run because "test" wasn't in the method name (and it was failing...

dynamic-shape gather leads to illegal mhlo error, but is it actually illegal?

4

### What happened? I'm working with targeting MHLO & IREE dynamic shapes from JAX, and in particular currently working on generating gather operations involving dynamic shapes, compiling those programs with...

FAILED tests/host_callback_test.py::HostCallbackTapTest::test_tap_remat8 - AssertionError: Found FAILED tests/host_callback_test.py::HostCallbackTapTest::test_tap_remat9 - AssertionError: Found FAILED tests/host_callback_test.py::HostCallbackTapTest::test_tap_transforms_doc - AssertionError: Found FAILED tests/host_callback_test.py::HostCallbackTapTest::test_tap_remat4 - AssertionError: Found

improve IndexError error message from abstracted values

2

This message could be improved to point to the specific reason (jit, pmap, fori_loop, etc) and user source line which caused the problematic tracing: "IndexError: Array slice indices must have...