jax

b9b2d4f9 - [pallas:mosaic] Recover `memory_space` from the aval in `aval_to_ir_type`

[pallas:mosaic] Recover `memory_space` from the aval in `aval_to_ir_type` Prior to this change `aval_to_ir_type` sometimes produced nonsensical memrefs, e.g. `memref<!tpu.semaphore, #tpu.memory_space<vmem>>`, because it failed to inspect the `memory_space` associated with the aval. PiperOrigin-RevId: 846173986