jax

2fd83c70 - [Mosaic GPU] Correctly propagate `approx=True` to `FragmentedArray.ex2` in `FragmentedArray.exp(approx=True)`.

[Mosaic GPU] Correctly propagate `approx=True` to `FragmentedArray.ex2` in `FragmentedArray.exp(approx=True)`. Harden the math tests a little in order to help avoid this kind of issues in the future. Note that the added test wouldn't have caught this issue exactly; it is hard to write a reliable test without pattern-matching the IR too closely. PiperOrigin-RevId: 934522639