Fix uniform distribution log_prob support

[juanitorduz]

Closes #2093

[fehiepsi]

This is actually a feature (for performance), not an issue. It is expected to enable validation like in @Qazalbash comment #2093 (comment) But maybe jax has improved a lot since then.

[juanitorduz]

Got it! So shall we merge this change or bring better docs on validation (as in the comment link above) ?

[fehiepsi]

I think we can enable validation by default:

numpyro/numpyro/distributions/distribution.py

Line 56 in 47335b8

_VALIDATION_ENABLED = False

[Qazalbash]

I think we can enable validation by default:

I second this. As a beginner, I had to deal with the same issue.

[juanitorduz]

Ok! I'll push the change :) Thanks for the input!

[juanitorduz]

@fehiepsi in 1663ec9 I reverted the change and simply enabled the validation (I left the tests). Let me know what you think about it :)

[juanitorduz]

mmm we are having now many failing tests:

AILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape0-Levy-<lambda>-params173] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape0-Levy-<lambda>-params174] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape0-_IntervalCensoredPoisson-None-params209] - UserWarning: For interval-censored, doubly-censored, or exact observations,lower bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape0-_IntervalCensoredPoisson-None-params211] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape0-_IntervalCensoredPoisson-None-params212] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape1-Levy-<lambda>-params173] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape1-Levy-<lambda>-params174] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape1-_IntervalCensoredPoisson-None-params209] - UserWarning: For interval-censored, doubly-censored, or exact observations,lower bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape1-_IntervalCensoredPoisson-None-params211] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape0-prepend_shape1-_IntervalCensoredPoisson-None-params212] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape0-Levy-<lambda>-params173] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape0-Levy-<lambda>-params174] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape0-_IntervalCensoredPoisson-None-params209] - UserWarning: For interval-censored, doubly-censored, or exact observations,lower bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape0-_IntervalCensoredPoisson-None-params211] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape0-_IntervalCensoredPoisson-None-params212] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape1-Levy-<lambda>-params173] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape1-Levy-<lambda>-params174] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape1-_IntervalCensoredPoisson-None-params209] - UserWarning: For interval-censored, doubly-censored, or exact observations,lower bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape1-_IntervalCensoredPoisson-None-params211] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_expand[sample_shape1-prepend_shape1-_IntervalCensoredPoisson-None-params212] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_dist_pytree[Levy-<lambda>-params173] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_dist_pytree[Levy-<lambda>-params174] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_distributions.py::test_dist_pytree[_IntervalCensoredPoisson-None-params209] - UserWarning: For interval-censored, doubly-censored, or exact observations,lower bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_dist_pytree[_IntervalCensoredPoisson-None-params211] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_dist_pytree[_IntervalCensoredPoisson-None-params212] - UserWarning: For left-censored observations, upper bound should be within the support of base_dist.
FAILED test/test_distributions.py::test_explicit_validate_args - ValueError: Normal distribution got invalid scale parameter.
FAILED test/test_distributions.py::test_interval_censored_logprob_edge_cases[1.0-1.0] - UserWarning: For interval-censored and doubly-censored observations,upper bound should greater than lower bound.
FAILED test/test_distributions.py::test_uniform_log_prob_outside_support - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_boundaries[0.0-1.0] - AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0

x and y -inf location mismatch:
 x: array(-0., dtype=float32)
 y: array(-inf)
FAILED test/test_distributions.py::test_uniform_log_prob_boundaries[-2.0-3.0] - AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0

x and y -inf location mismatch:
 x: array(-1.609438, dtype=float32)
 y: array(-inf)
FAILED test/test_distributions.py::test_uniform_log_prob_boundaries[1.0-5.0] - AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0

x and y -inf location mismatch:
 x: array(-1.386294, dtype=float32)
 y: array(-inf)
FAILED test/test_distributions.py::test_uniform_log_prob_boundaries[-5.0--1.0] - AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0

x and y -inf location mismatch:
 x: array(-1.386294, dtype=float32)
 y: array(-inf)
FAILED test/test_distributions.py::test_uniform_log_prob_broadcasting[batch_shape1] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_broadcasting[batch_shape2] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_broadcasting[batch_shape3] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_value_broadcasting[value_shape0] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_value_broadcasting[value_shape1] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_value_broadcasting[value_shape2] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_value_broadcasting[value_shape3] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_gof.py::test_gof[Levy-<lambda>-params173] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_gof.py::test_gof[Levy-<lambda>-params174] - ValueError: Levy distribution got invalid loc parameter.
FAILED test/test_handlers.py::test_mask_inf - ValueError: Unit distribution got invalid log_factor parameter.
FAILED test/test_pickle.py::test_pickle_hmcecs - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_util.py::test_format_shapes - ValueError: Normal distribution got invalid scale parameter.
= 44 failed, 7041 passed, 2082 skipped, 47 deselected, 39 xfailed, 61 xpassed in 1483.04s (0:24:43) =

shal we revert the validation change and keep this PR with the simple uniform fix?

[fehiepsi]

I think it's better to fix those bugs

• Levy loc constraint should be real:

  numpyro/numpyro/distributions/continuous.py

  Line 3282 in 3aed387

  "loc": constraints.positive,

• add validate_args=False to

  numpyro/test/test_distributions.py

  Line 3544 in 3aed387

  d = dist.Normal(0, 1)

• it's fine to set validate_args=False by default for Unit distribution

  numpyro/numpyro/distributions/distribution.py

  Line 1323 in 3aed387

  def __init__(self, log_factor: ArrayLike, *, validate_args: Optional[bool] = None):

• Similarly, let's set validate_args=False for censored distribution by default for now:

  numpyro/numpyro/distributions/censored.py

  Line 338 in 3aed387

  validate_args: Optional[bool] = None,

  I think its validation logic is not robust. cc @vanAmsterdam

[vanAmsterdam]

hmm I'm having trouble reproducing the errors thrown in test_exapand; I've tracked your fork branch @juanitorduz and tested in both 3.12 and 3.9 envs but no errors. any idea if i'm missing a step?

[juanitorduz]

@fehiepsi the CI has some tests failing

FAILED test/test_distributions.py::test_explicit_validate_args - ValueError: Normal distribution got invalid scale parameter.
FAILED test/test_distributions.py::test_uniform_log_prob_outside_support - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_boundaries[0.0-1.0] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_boundaries[-2.0-3.0] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_boundaries[1.0-5.0] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_boundaries[-5.0--1.0] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_broadcasting[batch_shape1] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_broadcasting[batch_shape2] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_broadcasting[batch_shape3] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_value_broadcasting[value_shape0] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_value_broadcasting[value_shape1] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_value_broadcasting[value_shape2] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_distributions.py::test_uniform_log_prob_value_broadcasting[value_shape3] - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_pickle.py::test_pickle_hmcecs - UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
FAILED test/test_util.py::test_format_shapes - ValueError: Normal distribution got invalid scale parameter.

However, I cannot reproduce locally (even with a brand-new Python env).

It's weird that the new tests are failing, given that we added the change to the log_prob. Any thoughts on why this could happen? (or how to debug it?) Thanks 🙏

[fehiepsi]

It seems there are issues with the tests:

• scale should be positive, we can use LogNormal there

  numpyro/test/test_util.py

  Line 115 in 212bccd

  scale = numpyro.sample("scale", dist.Normal(0, 1).expand([3]).to_event(1))

• for test_explicit_validate_args it is better so set validate_args=False like in my previous comment. validate_args(False) seems to do a different thing
• data in logistic_regression is wrong:

  numpyro/test/test_pickle.py

  Line 68 in 212bccd

  data = jnp.arange(10)

  you can replace it by an array with values in {0, 1}
• maybe revert the changes in uniform log_prob and fix the new tests?
• def enable_validation(is_validate: bool = False) -> None: should be True by default.

[juanitorduz]

Ok! the CI is 🟢 🙌 @fehiepsi any other comments or feedback :)

[fehiepsi]

Lgtm thanks!