jax

ced28167 - [Mosaic GPU] Use explicit recursion in rules instead of doing it automatically

[Mosaic GPU] Use explicit recursion in rules instead of doing it automatically Control-flow ops that have vector inputs or outputs will need to be specially adjusted. PiperOrigin-RevId: 730922072