Fix enable_validation

[renecotyfanboy]

Hi there,
Here is a super small fix for validation. Currently, the validation context is handled using _VALIDATION_ENABLED global variable. By default _VALIDATION_ENABLED=True and _validate_args=False by default, which trigger weird behaviors when enabling / disabling the validation through contexts. But since everything is handled through mutating _VALIDATION_ENABLED, it sometimes produces weird behaviors.

The simplest repro is :

print(dist.Uniform(0, 1).log_prob(-0.5)) # -0.0 (expected)

with numpyro.validation_enabled(True):
    print(dist.Uniform(0, 1).log_prob(-0.5)) # -inf (expected)

print(dist.Uniform(0, 1).log_prob(-0.5)) # -inf (expect -0.0)

The current fix produces the expected behavior.

[juanitorduz]

@fehiepsi : shall we keep _VALIDATION_ENABLED and _validate_args equal to True?

[fehiepsi]

We decided to enable validation by default. I think setting the attribute to the global one might fix the issue.

[juanitorduz]

Minor point, right @Qazalbash ?

[juanitorduz]

Based on the comment above, we should keep _VALIDATION_ENABLED = True

[renecotyfanboy]

This is fixed. Note that this changes the behavior by default and might be a breaking change for some codebases. However, I find that it is much better to have validation enabled by default, as it was quite unintuitive in the first place.

[Qazalbash]

@juanitorduz Yes, the global variable value should be reverted.

[Qazalbash]

Thanks @renecotyfanboy

[Qazalbash]

I think CI will trigger after @juanitorduz approves. Let's merge this PR only when the latest CI passes.

[juanitorduz]

@renecotyfanboy thanks! Can you push an empty commit (e.g., git commit -m"empty" --allow-empty ) to try to trigger the CI? 🙏

[renecotyfanboy]

Enabling validation by default broke some tests in the CI. Should I simply wrap them around with numpyro.validation_enabled(False) ? I am not sure if I am willing to fix them properly 😆