jax

3fc4ba29 - [Pallas TPU] Add lowerings for `lax.population_count_p` and `lax.clz_p`