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 😆

[Qazalbash]

Let's leave things as they were. What do you say @juanitorduz?

[fehiepsi]

Could we fix the tests instead? This indicates that we might not use correct params for those tests.

[juanitorduz]

Indeed, that is a better approach! I could help fix these tests if you feel tuck @renecotyfanboy :)

[renecotyfanboy]

I started fixing a bunch of tests here.

• Disabled validation for the categorical variables since most of them are not normalized in the tests
• Disabled validation for the Gaussian process contributed tests
• Used better MCMC params to ensure convergence for the functional map tests
• Used better random params for the negative binomial log prob tests

However, I see that it is taking a super long time for me to look at each case, each distribution, and find what is ill-defined in the parameter used in space or in the batch / event shape mismatch. If I am to continue fixing the test, I will probably switch to LLMs to iterate faster.

[juanitorduz]

Thank you @renecotyfanboy ! I think Claude (or similar) can be super useful here! If you need some help please let me know :)