jax

2ade7e75 - [pallas] Move the hardware_generation query in the code path that needs it

[pallas] Move the hardware_generation query in the code path that needs it This change allows us to lower and export Pallas calls even on machines that do not have TPUs, in many cases. PiperOrigin-RevId: 641841079