Possible optimization opportunity

Parse time increases with the number of instructions in the IR. There are some operations in the IR that could be cached. Instead of generating the same operation twice, just use the result of its value where appropriate. The functions safe_cast and extract_scalar in jax_primitives.py may be cached if they are used within the same function. It is unclear how valuable this caching may be since applying cse in MLIR should get rid of these. However, if there are many operations that can be cse'd directly generated by the MLIR code-gen in Python, it may be advantageous to cache these.

Start in jax_primitives.py by reading safe_cast and extract_scalar and tracing their call sites within a single function. Measure parse time and inspect whether repeated operations are already removed by MLIR CSE; done means establishing whether caching provides a meaningful benefit and defining the affected scope.