jax

7415c23e - [Mosaic GPU][NFC] Do not capture op references in the Scratch manager.

[Mosaic GPU][NFC] Do not capture op references in the Scratch manager. Before this change we created the `alloca` and `load` ops in `Scratch._create_ops_if_none` and also stored the pointers to these ops in the `Scratch` object. This happened when the first `async_copy` is called. Then later, at the end of the lowering, after we know all the SMEM we would need, we call `Scratch.finalize_size` to modify the alloc and load ops to reflect the real size. As long as there was no `canonicalize` called in between, the pointers remained valid and the above worked fine. However, now I'm implementing support for `inline_mgpu` in Warp-group semantics and we have a new situation. `async_copy` within the `inline_mgpu` block is called early on in the `Pallas` -> `WG` lowering. This creates the `alloca` and `load` ops. Then later, we call `canonicalize`, and afterwards proceed with the `WG` -> `Lane` lowering. Finally, we call `Scratch.finalize_size`, but now the stored pointers to `alloca` and `load` are dangling. To avoid this, this changes always traverses the IR to find the right ops and doesn't store anything. PiperOrigin-RevId: 781934910