code review suggestions

pyrenew/convolve.py

@@ -82,7 +82,7 @@ def _new_scanner(
             * jnp.einsum("i...,i...->...", array_to_convolve, history_subset)
         )
         latest = jnp.concatenate(
-            [history_subset[1:], jnp.expand_dims(new_val, axis=0)], axis=0
+            [history_subset[1:], new_val[jnp.newaxis]], axis=0
         )
         return latest, new_val
 
@@ -166,7 +166,7 @@ def _new_scanner(
             m2 * m_net1 * jnp.einsum("i...,i...->...", arr2, history_subset)
         )
         latest = jnp.concatenate(
-            [history_subset[1:], jnp.expand_dims(new_val, axis=0)], axis=0
+            [history_subset[1:], new_val[jnp.newaxis]], axis=0
         )
         return latest, (new_val, m_net1)
 

pyrenew/latent/infections.py

@@ -88,11 +88,11 @@ def sample(
                 f"generation interval length: {gen_int.size}."
             )
 
-        if I0.ndim != Rt.ndim:
+        if I0.shape[1:] != Rt.shape[1:]:
             raise ValueError(
-                "Initial infections and Rt must have the same dimensions. "
-                f"Got initial infections of shape {I0.shape} "
-                f"and Rt of shape {Rt.shape}."
+                "Initial infections and Rt must have the same batch shapes. "
+                f"Got initial infections of batch shape {I0.shape[1:]} "
+                f"and Rt of batch shape {Rt.shape[1:]}."
             )
 
         gen_int_rev = jnp.flip(gen_int)

pyrenew/latent/infectionswithfeedback.py

@@ -150,11 +150,11 @@ def sample(
                 f"and generation interval length {gen_int.size}."
             )
 
-        if I0.ndim != Rt.ndim:
+        if I0.shape[1:] != Rt.shape[1:]:
             raise ValueError(
-                "Initial infections and Rt must have the same dimensions. "
-                f"Got initial infections of shape {I0.shape} "
-                f"and Rt of shape {Rt.shape}."
+                "Initial infections and Rt must have the same batch shapes. "
+                f"Got initial infections of batch shape {I0.shape[1:]} "
+                f"and Rt of batch shape {Rt.shape[1:]}."
             )
 
         gen_int_rev = jnp.flip(gen_int)

test/test_infection_and_infectionwithfeedback.py

@@ -60,7 +60,7 @@ def test_infections_with_feedback_invalid_inputs():
 
         with pytest.raises(
             ValueError,
-            match="Initial infections and Rt must have the same dimensions.",
+            match="Initial infections and Rt must have the same batch shapes.",
         ):
             InfectionsWithFeedback(
                 gen_int=gen_int,
@@ -70,7 +70,7 @@ def test_infections_with_feedback_invalid_inputs():
 
         with pytest.raises(
             ValueError,
-            match="Initial infections and Rt must have the same dimensions.",
+            match="Initial infections and Rt must have the same batch shapes.",
         ):
             infections(
                 gen_int=gen_int,