Bump to jax 0.9.1

pyrenew/math.py

@@ -27,6 +27,7 @@ def _positive_ints_like(vec: ArrayLike) -> jnp.ndarray:
     jnp.ndarray
         The resulting array ``[1, ..., n]``.
     """
+    vec = jnp.asarray(vec)
     return jnp.arange(1, jnp.size(vec) + 1)
 
 

test/test_convolve.py

@@ -82,7 +82,7 @@ def transform_a(x: any):
         [
             jnp.ones(3),
             jnp.array(np.array([0.5, 0.3, 0.2] * 3)).reshape(3, 3),
-            jnp.ones((3, 3)),
+            jnp.array([[-0.25, 0.25, 0.25], [0, 1, 0], [0.5, -0.5, 0.5]]),
             t.ExpTransform(),
         ],
     ],
@@ -95,14 +95,13 @@ def test_convolve_scanner_using_scan(arr, history, multipliers, transform):
     """
     scanner = pc.new_convolve_scanner(arr, transform)
 
-    _, result = jax.lax.scan(f=scanner, init=history, xs=multipliers)
+    _, result = jax.lax.scan(f=scanner, init=jnp.array(history), xs=multipliers)
 
     result_not_scanned = []
     for multiplier in multipliers:
         history, new_val = scanner(history, multiplier)
         result_not_scanned.append(new_val)
-
-    assert jnp.array_equal(result, result_not_scanned)
+    assert_array_equal(result, result_not_scanned)
 
 
 @pytest.mark.parametrize(

test/test_infectionwithfeedback_plate_compatibility.py

@@ -20,7 +20,9 @@ def test_infections_with_feedback_plate_compatibility():
     Rt = jnp.ones((10, 5))
     gen_int = jnp.array([0.4, 0.25, 0.25, 0.1])
 
-    inf_feed_strength = DistributionalVariable("inf_feed_strength", dist.Beta(1, 1))
+    inf_feed_strength = DistributionalVariable(
+        "inf_feed_strength", dist.LogNormal(0.0, 1.0)
+    )
     inf_feedback_pmf = DeterministicPMF(name="inf_feedback_pmf", value=gen_int)
 
     # Test the InfectionsWithFeedback class

test/test_observation_counts.py

@@ -150,7 +150,7 @@ class TestCountsWithPriors:
     def test_with_stochastic_ascertainment(self, short_shedding_pmf):
         """Test with uncertain ascertainment rate parameter."""
         delay = DeterministicPMF("delay", jnp.array([0.2, 0.5, 0.3]))
-        ascertainment = DistributionalVariable("ihr", dist.Beta(2, 100))
+        ascertainment = DistributionalVariable("ihr", dist.LogNormal(-1, 100.0))
         concentration = DeterministicVariable("conc", 10.0)
 
         process = Counts(