Allow registering of custom exception handlers for potential_fn computations

[Balandat]

In some cases evaluation the potential funciton may result in numerical issues. Currently the code hard-codes the handling of a RuntimeError raised when matrices are (numerically) singular. This PR adds the ability to register custom exception handlers. This allows other code depending on pyro to register custom exception handlers without having to modify core pyro code.

There are some other places in which potential_fn is called that could benefit from being guarded by these handlers (one is HMC._find_reasonable_step_size). I'm not sure what the right thing to do there is when encountering numerical isssues, but happy to add this in as needed.

[Balandat]

@fritzo is this (or something like this) something you'd be willing to merge in? This will be quite useful for folks who may use some more numerically challenging setups (e.g. GPs) that rely on implementations that throw some other errors.

cc @dme65, @saitcakmak

[fehiepsi]

Looks great to me. Thanks, @Balandat!

The utility will be exposed in docs around this path: https://docs.pyro.ai/en/stable/ops.html#pyro.ops.integrator.potential_grad

[fehiepsi]

Could you fix the lint issue?

[Balandat]

Sure, will do.

Any thoughts how to handle numerical issues of this kind in the following:

pyro/pyro/infer/mcmc/hmc.py

Line 176 in 0b1818c

potential_energy = self.potential_fn(z)

[fehiepsi]

I think you can do similarly

try:
    potential_energy = self.potential_fn(z)
except ...:
    # skip finding reasonable step size
    return step_size

[fritzo]

This seems reasonable. Since we may also want to use this or something like it in NumPyro, could we discuss the design advantages of the proposed solution over alternative solutions?

1. A higher-level solution: In case of error, MH reject the trajectory and sample a different velocity. I'm unsure this makes sense, but it also seems unclear whether error-as-zero-grad makes sense 🤷

2. A lower-level solution: Wrap the potential function outside of MCMC to silence exceptions:

def safe_potential(fn: Callable[[dict], torch.Tensor]) -> Callable[[dict], torch.Tensor]:
    def safe_fn(z: dict) -> torch.Tensor:
        try:
            return fn(z)
        except RuntimeError as e:
            if "singular" in str(e) or "input is not positive-definite" in str(e):
                return next(iter(z.values())).new_tensor(math.nan)
            raise e
    return safe_fn

which I guess would also require a change to potential_grad()

- grads = grad(potential_energy, z_nodes)
+ grads = grad(potential_energy, z_nodes, allow_unused=True)

[fehiepsi]

Currently, when errors happen, we set potential_energy to nan, which will lead to the rejection of the proposal. grad can be zeros, ones, or empty - IIUC it won't affect the logic. In NumPyro, jax won't throw errors for numerical computations - rather than that, it throws some invalid outputs like nan or inf, so probably we don't need those exceptions.

[Balandat]

This seems reasonable. Since we may also want to use this or something like it in NumPyro, could we discuss the design advantages of the proposed solution over alternative solutions?

I'm not an expert on what the proper thing to do algorithmically is. Ideally the handling of this can happen at a higher level as part of the algorithm. If the proposal is rejected by setting things to nan and this is the right thing to do, then great.

The downside of the lower-level solution you proposed is that it will require the user to understand the details of the algorithm to implement this kind of wrapper, rather than "I want to handle these kinds of exceptions and have the algorithm deal with properly rejecting the proposal".

[fritzo]

@Balandat sorry linting seems to be a hassle. We usually use the make targets make format and make lint to do all the linting and formatting automatically.

[Balandat]

@Balandat sorry linting seems to be a hassle. We usually use the make targets make format and make lint to do all the linting and formatting automatically.

No worries - I should have run make format. Note though that make lint currently spits out hundreds if not thousands of errors in files unrelated to this PR.

[fehiepsi]

LGTM