jax

44be67a6 - Fix test failure in newly added test.

Fix test failure in newly added test. It and other loop tests were badly located in PallasCallPoisonTest, which seems like the wrong place, and had sizes that didn't work on TPUv3. PiperOrigin-RevId: 929162162