jax

dfcfb360 - Pallas GPU no longer falls back to lax.pow for integer powers

Pallas GPU no longer falls back to lax.pow for integer powers Instead the lowering computes the power in a loop by squaring, similarly to how we do it in the StableHLO lowering. Fixes #21928. PiperOrigin-RevId: 644313113