Possible optimization opportunity

[erick-xanadu]

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.