Enforce Array consistency across distributions

distreqx/distributions/_bernoulli.py

@@ -1,6 +1,6 @@
 """Bernoulli distribution."""
 
-from typing import Optional, Union
+from typing import Optional
 
 import jax
 import jax.numpy as jnp
@@ -24,13 +24,13 @@ class Bernoulli(
     Bernoulli distribution with parameter `probs`, the probability of outcome `1`.
     """
 
-    _logits: Union[Array, None]
-    _probs: Union[Array, None]
+    _logits: Array | None
+    _probs: Array | None
 
     def __init__(
         self,
-        logits: Optional[Array] = None,
-        probs: Optional[Array] = None,
+        logits: Optional[float | Array] = None,
+        probs: Optional[float | Array] = None,
     ):
         """Initializes a Bernoulli distribution.
 
@@ -48,11 +48,9 @@ def __init__(
                 f"One and exactly one of `logits` and `probs` should be `None`, "
                 f"but `logits` is {logits} and `probs` is {probs}."
             )
-        if (not isinstance(logits, jax.Array)) and (not isinstance(probs, jax.Array)):
-            raise ValueError("`logits` and `probs` are not jax arrays.")
         # Parameters of the distribution.
-        self._probs = None if probs is None else probs
-        self._logits = None if logits is None else logits
+        self._probs = None if probs is None else jnp.asarray(probs)
+        self._logits = None if logits is None else jnp.asarray(logits)
 
     @property
     def logits(self) -> Array:

distreqx/distributions/_categorical.py

@@ -44,8 +44,6 @@ def __init__(self, logits: Optional[Array] = None, probs: Optional[Array] = None
                 f"One and exactly one of `logits` and `probs` should be `None`, "
                 f"but `logits` is {logits} and `probs` is {probs}."
             )
-        if (not isinstance(logits, jax.Array)) and (not isinstance(probs, jax.Array)):
-            raise ValueError("`logits` and `probs` are not jax arrays.")
 
         self._probs = None if probs is None else normalize(probs=probs)
         self._logits = None if logits is None else normalize(logits=logits)

distreqx/distributions/_mvn_diag.py

@@ -1,6 +1,6 @@
 """MultivariateNormalDiag distribution."""
 
-from typing import Optional
+from typing import Optional, Union
 
 import equinox as eqx
 import jax
@@ -53,7 +53,11 @@ class MultivariateNormalDiag(AbstractMultivariateNormalFromBijector):
     bijector: AbstractBijector
     scale_diag: Array
 
-    def __init__(self, loc: Optional[Array] = None, scale_diag: Optional[Array] = None):
+    def __init__(
+        self,
+        loc: Optional[Union[float, Array]] = None,
+        scale_diag: Optional[Union[float, Array]] = None,
+    ):
         """Initializes a MultivariateNormalDiag distribution.
 
         **Arguments:**
@@ -63,6 +67,8 @@ def __init__(self, loc: Optional[Array] = None, scale_diag: Optional[Array] = No
         - `scale_diag`: Vector of standard deviations.  If not specified, it
             defaults to ones. At least one of `loc` and`scale_diag` must be specified.
         """
+        loc = None if loc is None else jnp.asarray(loc)
+        scale_diag = None if scale_diag is None else jnp.asarray(scale_diag)
         _check_parameters(loc, scale_diag)
 
         if scale_diag is None and loc is not None:

distreqx/distributions/_mvn_from_bijector.py

@@ -1,6 +1,7 @@
 """MultivariateNormalFromBijector distribution."""
 
 from collections.abc import Callable
+from typing import Union
 
 import equinox as eqx
 import jax
@@ -116,7 +117,7 @@ class MultivariateNormalFromBijector(AbstractMultivariateNormalFromBijector):
     distribution: AbstractDistribution
     bijector: AbstractBijector
 
-    def __init__(self, loc: Array, scale: AbstractLinearBijector):
+    def __init__(self, loc: Union[float, Array], scale: AbstractLinearBijector):
         """Initializes the distribution.
 
         **Arguments:**
@@ -125,6 +126,7 @@ def __init__(self, loc: Array, scale: AbstractLinearBijector):
         - `scale`: The bijector specifying the linear transformation `A @ z`, as
             described in the class docstring.
         """
+        loc = jnp.asarray(loc)
         _check_input_parameters_are_valid(scale, loc)
 
         # Build a standard multivariate Gaussian.

distreqx/distributions/_mvn_full_covariance.py

@@ -1,6 +1,6 @@
 """MultivariateNormalFullCovariance distribution."""
 
-from typing import Optional
+from typing import Optional, Union
 
 import equinox as eqx
 import jax.numpy as jnp
@@ -80,8 +80,8 @@ class MultivariateNormalFullCovariance(
 
     def __init__(
         self,
-        loc: Optional[Array] = None,
-        covariance_matrix: Optional[Array] = None,
+        loc: Optional[Union[float, Array]] = None,
+        covariance_matrix: Optional[Union[float, Array]] = None,
     ):
         """Initializes a MultivariateNormalFullCovariance distribution.
 
@@ -93,6 +93,10 @@ def __init__(
             symmetric positive definite matrix. If not specified, it defaults
             to the identity matrix.
         """
+        loc = None if loc is None else jnp.asarray(loc)
+        covariance_matrix = (
+            None if covariance_matrix is None else jnp.asarray(covariance_matrix)
+        )
         _check_parameters(loc, covariance_matrix)
 
         if loc is not None:

distreqx/distributions/_mvn_tri.py

@@ -1,6 +1,6 @@
 """MultivariateNormalTri distribution."""
 
-from typing import Optional
+from typing import Optional, Union
 
 import equinox as eqx
 import jax
@@ -73,8 +73,8 @@ class MultivariateNormalTri(AbstractMultivariateNormalFromBijector):
 
     def __init__(
         self,
-        loc: Optional[Array] = None,
-        scale_tri: Optional[Array] = None,
+        loc: Optional[Union[float, Array]] = None,
+        scale_tri: Optional[Union[float, Array]] = None,
         is_lower: bool = True,
     ):
         """Initializes a MultivariateNormalTri distribution.
@@ -93,6 +93,8 @@ def __init__(
         - `is_lower`: Indicates if `scale_tri` is lower (if True) or upper (if False)
             triangular.
         """
+        loc = None if loc is None else jnp.asarray(loc)
+        scale_tri = None if scale_tri is None else jnp.asarray(scale_tri)
         _check_parameters(loc, scale_tri)
 
         if loc is not None:

distreqx/distributions/_one_hot_categorical.py

@@ -39,8 +39,6 @@ def __init__(self, logits: Optional[Array] = None, probs: Optional[Array] = None
                 f"One and exactly one of `logits` and `probs` should be `None`, "
                 f"but `logits` is {logits} and `probs` is {probs}."
             )
-        if (not isinstance(logits, jax.Array)) and (not isinstance(probs, jax.Array)):
-            raise ValueError("`logits` and `probs` are not jax arrays.")
 
         self._probs = None if probs is None else normalize(probs=probs)
         self._logits = None if logits is None else normalize(logits=logits)

distreqx/distributions/_uniform.py

@@ -1,5 +1,6 @@
 """Uniform distribution."""
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Float, Key
@@ -13,8 +14,8 @@ class Uniform(
 ):
     """Uniform distribution with `low` and `high` parameters."""
 
-    low: Float[Array, "..."]
-    high: Float[Array, "..."]
+    low: Float[Array, "..."] = eqx.field(converter=jnp.asarray)
+    high: Float[Array, "..."] = eqx.field(converter=jnp.asarray)
 
     @property
     def range(self) -> Array:

tests/bernoulli_test.py

@@ -50,6 +50,13 @@ def test_raises_on_invalid_inputs(self, name, dist_params):
         with self.assertRaises(ValueError):
             self.dist(**dist_params)
 
+    def test_accepts_raw_python_floats(self):
+        """Bernoulli should cast plain Python floats to arrays, like `Normal` does."""
+        dist = self.dist(probs=0.3)
+        self.assertIsInstance(dist.probs, jax.Array)
+        dist = self.dist(logits=0.5)
+        self.assertIsInstance(dist.logits, jax.Array)
+
     @parameterized.expand(
         [
             ("1d probs, 1-tuple shape", {"probs": [0.1, 0.5, 0.3]}, (3,)),

tests/mvn_diag_test.py

@@ -30,6 +30,14 @@ def test_invalid_parameters(self):
             dist_kwargs={"loc": jnp.zeros((3, 5)), "scale_diag": jnp.ones((3, 4))}
         )
 
+    def test_raw_python_float_raises_value_error_not_attribute_error(self):
+        """A raw (uncast) `loc`/`scale_diag` should fail shape validation cleanly,
+        instead of crashing on a missing `.shape` attribute."""
+        with self.assertRaises(ValueError):
+            MultivariateNormalDiag(loc=1.0, scale_diag=None)
+        with self.assertRaises(ValueError):
+            MultivariateNormalDiag(loc=None, scale_diag=1.0)
+
     @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)])
     def test_sample_dtype(self, name, dtype):
         dist_params = {

tests/mvn_from_bijector_test.py

@@ -58,6 +58,13 @@ def test_raises_on_wrong_inputs(self, name, event_dims, loc):
         with self.assertRaises(ValueError):
             MultivariateNormalFromBijector(loc, bij)
 
+    def test_raw_python_float_raises_value_error_not_attribute_error(self):
+        """A raw (uncast) `loc` should fail shape validation cleanly, instead of
+        crashing on a missing `.ndim` attribute."""
+        bij = MockLinear(4)
+        with self.assertRaises(ValueError):
+            MultivariateNormalFromBijector(1.0, bij)
+
     @parameterized.expand([("no broadcast", jnp.ones((4,)), jnp.zeros((4,)), (4,))])
     def test_loc_scale_and_shapes(self, name, diag, loc, expected_shape):
         scale = DiagLinear(diag)

tests/mvn_full_covariance_test.py

@@ -55,6 +55,14 @@ def test_invalid_parameters(self):
             dist_kwargs={"loc": jnp.zeros((5,)), "covariance_matrix": jnp.eye(4)}
         )
 
+    def test_raw_python_float_raises_value_error_not_attribute_error(self):
+        """A raw (uncast) `loc`/`covariance_matrix` should fail shape validation
+        cleanly, instead of crashing on a missing `.shape`/`.ndim` attribute."""
+        with self.assertRaises(ValueError):
+            MultivariateNormalFullCovariance(loc=1.0, covariance_matrix=None)
+        with self.assertRaises(ValueError):
+            MultivariateNormalFullCovariance(loc=None, covariance_matrix=1.0)
+
     @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)])
     def test_sample_dtype(self, name, dtype):
         dist_params = {

tests/mvn_tri_test.py

@@ -39,6 +39,14 @@ def test_invalid_parameters(self):
             dist_kwargs={"loc": jnp.zeros((5,)), "scale_tri": jnp.ones((4, 4))}
         )
 
+    def test_raw_python_float_raises_value_error_not_attribute_error(self):
+        """A raw (uncast) `loc`/`scale_tri` should fail shape validation cleanly,
+        instead of crashing on a missing `.shape`/`.ndim` attribute."""
+        with self.assertRaises(ValueError):
+            MultivariateNormalTri(loc=1.0, scale_tri=None)
+        with self.assertRaises(ValueError):
+            MultivariateNormalTri(loc=None, scale_tri=1.0)
+
     @parameterized.expand([("float32", jnp.float32), ("float64", jnp.float64)])
     def test_sample_dtype(self, name, dtype):
         dist_params = {

tests/uniform_test.py

@@ -150,6 +150,14 @@ def test_with_two_distributions(self, function_string):
         )
         self.assertEqual(result.shape, expected_shape)
 
+    def test_accepts_raw_python_floats(self):
+        """Uniform should cast plain Python floats to arrays, like `Normal` does."""
+        dist = Uniform(0.0, 1.0)
+        self.assertIsInstance(dist.low, jax.Array)
+        self.assertIsInstance(dist.high, jax.Array)
+        sample = dist.sample(self.key)
+        self.assertEqual(sample.shape, ())
+
     def test_jittable(self):
         @jax.jit
         def create_and_sample(key):