jax

66a3f87a - Rollback for: Implement initial vmap over pallas_call w/ ragged inputs (via jumbles)

Rollback for: Implement initial vmap over pallas_call w/ ragged inputs (via jumbles) It can cause issues in x32 when trying to get the aval for array dimension sizes that are larger than i32. Reverts 24394a1b03f01138219013f4773104b834e498b7 PiperOrigin-RevId: 664742891