Make cond test more robust (#1029)

Author: tcbegley

test/contrib/test_control_flow.py

@@ -190,14 +190,8 @@ def false_fun(_):
     x = mcmc.get_samples()["x"]
     assert x.shape == (10_000,)
     assert_allclose(
-        [
-            x.mean(),
-            x.std(),
-            x[x > 2.0].mean(),
-            x[x > 2.0].std(),
-            x[x < 2.0].mean(),
-            x[x < 2.0].std(),
-        ],
-        [2.0, jnp.sqrt(5.0), 4.0, 1.0, 0.0, 1.0],
+        [x[x > 2.0].mean(), x[x > 2.0].std(), x[x < 2.0].mean(), x[x < 2.0].std()],
+        [4.01, 0.965, -0.01, 0.965],
         atol=0.1,
     )
+    assert_allclose([x.mean(), x.std()], [2.0, jnp.sqrt(5.0)], atol=0.5)