feedback 1/n

Author: juanitorduz

numpyro/distributions/distribution.py

@@ -56,7 +56,7 @@
 _VALIDATION_ENABLED = True
 
 
-def enable_validation(is_validate: bool = False) -> None:
+def enable_validation(is_validate: bool = True) -> None:
     """
     Enable or disable validation checks in NumPyro. Validation checks provide useful warnings and
     errors, e.g. NaN checks, validating distribution arguments and support values, etc. which is

test/test_distributions.py

@@ -3693,13 +3693,13 @@ def test_vmap_validate_args():
 
 def test_explicit_validate_args():
     # Check validation passes for valid parameters.
-    d = dist.Normal(0, 1)
-    d.validate_args(False)
+    d = dist.Normal(0, 1, validate_args=False)
+    d.validate_args()
 
     # Check validation fails for invalid parameters.
     d = dist.Normal(0, -1)
     with pytest.raises(ValueError, match="got invalid scale parameter"):
-        d.validate_args(False)
+        d.validate_args()
 
     # Check validation is skipped for strict=False and raises an error for strict=True.
     jitted = jax.jit(

test/test_pickle.py

@@ -65,7 +65,7 @@ def bernoulli_model():
 
 
 def logistic_regression():
-    data = jnp.arange(10)
+    data = random.choice(random.PRNGKey(0), jnp.array([0, 1]), (10,))
     x = numpyro.sample("x", dist.Normal(0, 1))
     with numpyro.plate("N", 10, subsample_size=2):
         batch = numpyro.subsample(data, 0)

test/test_util.py

@@ -112,7 +112,7 @@ def test_format_shapes():
 
     def model_test():
         mean = numpyro.param("mean", jnp.zeros(len(data)))
-        scale = numpyro.sample("scale", dist.Normal(0, 1).expand([3]).to_event(1))
+        scale = numpyro.sample("scale", dist.LogNormal(0, 1).expand([3]).to_event(1))
         scale = scale.sum()
         with numpyro.plate("data", len(data), subsample_size=10) as ind:
             batch = data[ind]