diff --git a/distreqx/distributions/_bernoulli.py b/distreqx/distributions/_bernoulli.py index 01be5a8..c14f0f0 100644 --- a/distreqx/distributions/_bernoulli.py +++ b/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: diff --git a/distreqx/distributions/_categorical.py b/distreqx/distributions/_categorical.py index 54d284d..9bc8814 100644 --- a/distreqx/distributions/_categorical.py +++ b/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) diff --git a/distreqx/distributions/_mvn_diag.py b/distreqx/distributions/_mvn_diag.py index 5755a87..eb1f923 100644 --- a/distreqx/distributions/_mvn_diag.py +++ b/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: diff --git a/distreqx/distributions/_mvn_from_bijector.py b/distreqx/distributions/_mvn_from_bijector.py index 4ee3448..7a755f0 100644 --- a/distreqx/distributions/_mvn_from_bijector.py +++ b/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. diff --git a/distreqx/distributions/_mvn_full_covariance.py b/distreqx/distributions/_mvn_full_covariance.py index 235a939..e41e494 100644 --- a/distreqx/distributions/_mvn_full_covariance.py +++ b/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: diff --git a/distreqx/distributions/_mvn_tri.py b/distreqx/distributions/_mvn_tri.py index 2de7a6e..f6c07ae 100644 --- a/distreqx/distributions/_mvn_tri.py +++ b/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: diff --git a/distreqx/distributions/_one_hot_categorical.py b/distreqx/distributions/_one_hot_categorical.py index 9cc39ce..3e65df4 100644 --- a/distreqx/distributions/_one_hot_categorical.py +++ b/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) diff --git a/distreqx/distributions/_uniform.py b/distreqx/distributions/_uniform.py index d059bbb..e906243 100644 --- a/distreqx/distributions/_uniform.py +++ b/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: diff --git a/tests/bernoulli_test.py b/tests/bernoulli_test.py index a4bd8c1..6d74e9b 100644 --- a/tests/bernoulli_test.py +++ b/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,)), diff --git a/tests/mvn_diag_test.py b/tests/mvn_diag_test.py index b897e93..942249e 100644 --- a/tests/mvn_diag_test.py +++ b/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 = { diff --git a/tests/mvn_from_bijector_test.py b/tests/mvn_from_bijector_test.py index 734fc4e..1e00b1c 100644 --- a/tests/mvn_from_bijector_test.py +++ b/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) diff --git a/tests/mvn_full_covariance_test.py b/tests/mvn_full_covariance_test.py index abab03d..4d6f83b 100644 --- a/tests/mvn_full_covariance_test.py +++ b/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 = { diff --git a/tests/mvn_tri_test.py b/tests/mvn_tri_test.py index e397a9c..56d9d09 100644 --- a/tests/mvn_tri_test.py +++ b/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 = { diff --git a/tests/uniform_test.py b/tests/uniform_test.py index 94bec40..674817d 100644 --- a/tests/uniform_test.py +++ b/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):