From c1c1c08aec0b6ce866892a2d1ddeb522ac7e22cb Mon Sep 17 00:00:00 2001 From: Gary Allen Date: Mon, 6 Jul 2026 17:15:35 +0200 Subject: [PATCH 1/4] Cast distribution parameters to arrays consistently Uniform had no __init__ at all, so `low`/`high` were stored as whatever was passed in (e.g. bare Python floats) instead of being cast to jax arrays like Normal, Beta, Gamma, etc. already do. Bernoulli/Categorical/ OneHotCategorical actively rejected non-jax-array probs/logits via an isinstance check instead of casting them, and the four MultivariateNormal classes accessed .shape/.ndim on loc/scale/covariance before casting, crashing with AttributeError for plain floats or lists. Cast inputs to arrays up front in all of these, matching the existing convention, so construction behaves consistently across the library regardless of whether callers pass raw Python numbers/lists or arrays. --- distreqx/distributions/_bernoulli.py | 10 ++++------ distreqx/distributions/_categorical.py | 2 -- distreqx/distributions/_mvn_diag.py | 10 ++++++++-- distreqx/distributions/_mvn_from_bijector.py | 4 +++- distreqx/distributions/_mvn_full_covariance.py | 10 +++++++--- distreqx/distributions/_mvn_tri.py | 8 +++++--- distreqx/distributions/_one_hot_categorical.py | 2 -- distreqx/distributions/_uniform.py | 17 +++++++++++++++++ tests/bernoulli_test.py | 7 +++++++ tests/mvn_diag_test.py | 8 ++++++++ tests/mvn_from_bijector_test.py | 7 +++++++ tests/mvn_full_covariance_test.py | 8 ++++++++ tests/mvn_tri_test.py | 8 ++++++++ tests/uniform_test.py | 8 ++++++++ 14 files changed, 90 insertions(+), 19 deletions(-) diff --git a/distreqx/distributions/_bernoulli.py b/distreqx/distributions/_bernoulli.py index 01be5a8..b3c4403 100644 --- a/distreqx/distributions/_bernoulli.py +++ b/distreqx/distributions/_bernoulli.py @@ -29,8 +29,8 @@ class Bernoulli( def __init__( self, - logits: Optional[Array] = None, - probs: Optional[Array] = None, + logits: Optional[Union[float, Array]] = None, + probs: Optional[Union[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..88fb77f 100644 --- a/distreqx/distributions/_uniform.py +++ b/distreqx/distributions/_uniform.py @@ -1,5 +1,7 @@ """Uniform distribution.""" +from typing import Union + import jax import jax.numpy as jnp from jaxtyping import Array, Float, Key @@ -16,6 +18,21 @@ class Uniform( low: Float[Array, "..."] high: Float[Array, "..."] + def __init__( + self, + low: Union[float, Float[Array, "..."]], + high: Union[float, Float[Array, "..."]], + ): + """Initializes a Uniform distribution. + + **Arguments:** + + - `low`: Lower bound of the distribution. + - `high`: Upper bound of the distribution. + """ + self.low = jnp.asarray(low) + self.high = jnp.asarray(high) + @property def range(self) -> Array: return self.high - self.low 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): From 7c4c5abb9728fc362b34eaeacdec764ca0d4063e Mon Sep 17 00:00:00 2001 From: Gary Allen Date: Tue, 7 Jul 2026 12:35:16 +0200 Subject: [PATCH 2/4] Use | instead of Union and eqx.field instead of __init__ --- distreqx/distributions/_bernoulli.py | 10 +++++----- distreqx/distributions/_uniform.py | 20 +++----------------- 2 files changed, 8 insertions(+), 22 deletions(-) diff --git a/distreqx/distributions/_bernoulli.py b/distreqx/distributions/_bernoulli.py index b3c4403..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[Union[float, Array]] = None, - probs: Optional[Union[float, Array]] = None, + logits: Optional[float | Array] = None, + probs: Optional[float | Array] = None, ): """Initializes a Bernoulli distribution. diff --git a/distreqx/distributions/_uniform.py b/distreqx/distributions/_uniform.py index 88fb77f..3a8ca65 100644 --- a/distreqx/distributions/_uniform.py +++ b/distreqx/distributions/_uniform.py @@ -5,6 +5,7 @@ import jax import jax.numpy as jnp from jaxtyping import Array, Float, Key +import equinox as eqx from ._distribution import AbstractSTDDistribution, AbstractSurvivalDistribution @@ -15,23 +16,8 @@ class Uniform( ): """Uniform distribution with `low` and `high` parameters.""" - low: Float[Array, "..."] - high: Float[Array, "..."] - - def __init__( - self, - low: Union[float, Float[Array, "..."]], - high: Union[float, Float[Array, "..."]], - ): - """Initializes a Uniform distribution. - - **Arguments:** - - - `low`: Lower bound of the distribution. - - `high`: Upper bound of the distribution. - """ - self.low = jnp.asarray(low) - self.high = jnp.asarray(high) + low: Float[Array, "..."] = eqx.field(converter=jnp.asarray) + high: Float[Array, "..."] = eqx.field(converter=jnp.asarray) @property def range(self) -> Array: From bc8590eb20534be82a5bccdc73cbb67baa7320f5 Mon Sep 17 00:00:00 2001 From: Gary Allen Date: Tue, 7 Jul 2026 12:41:02 +0200 Subject: [PATCH 3/4] Removed unused import --- distreqx/distributions/_uniform.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/distreqx/distributions/_uniform.py b/distreqx/distributions/_uniform.py index 3a8ca65..dffa733 100644 --- a/distreqx/distributions/_uniform.py +++ b/distreqx/distributions/_uniform.py @@ -1,7 +1,5 @@ """Uniform distribution.""" -from typing import Union - import jax import jax.numpy as jnp from jaxtyping import Array, Float, Key From 6fed5198a85fd5e392c3865c3eb45610db2f705b Mon Sep 17 00:00:00 2001 From: Gary Allen Date: Tue, 7 Jul 2026 12:44:49 +0200 Subject: [PATCH 4/4] Ran pre-commit --- distreqx/distributions/_uniform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/distreqx/distributions/_uniform.py b/distreqx/distributions/_uniform.py index dffa733..e906243 100644 --- a/distreqx/distributions/_uniform.py +++ b/distreqx/distributions/_uniform.py @@ -1,9 +1,9 @@ """Uniform distribution.""" +import equinox as eqx import jax import jax.numpy as jnp from jaxtyping import Array, Float, Key -import equinox as eqx from ._distribution import AbstractSTDDistribution, AbstractSurvivalDistribution