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?