jax

f600caa2 - Use @register_lowering to register Pallas GPU lowering rules

Use @register_lowering to register Pallas GPU lowering rules This leads to slightly more compact code, but should otherwise be identical. PiperOrigin-RevId: 635442002