jax version of GeneralizedExtremeValue distribution has wrong support for concentration=0

[jobrachem]

It seems like the support of the GEV is (-1.0, inf) for concentration=0, rather than the expected (-inf, inf). I am using:

• tensorflow-probability 0.24.0
• jax 0.4.30

I only tested this for the tensorflow_probability.substrates.jax.distributions version.

Otherwise, the support seems to be fine. Here's a little experiment for reproducing my findings:

import tensorflow_probability.substrates.jax.distributions as tfd
import jax.numpy as jnp
import pandas as pd
from itertools import product

def nominal_gev_support(loc, scale, concentration):
    if jnp.allclose(concentration, 0.0):
        return (-jnp.inf, jnp.inf)
    
    if concentration > 0.0:
        return (loc - scale / concentration, jnp.inf)
    
    if concentration < 0.0:
        return (-jnp.inf, loc - scale / concentration)

x = jnp.linspace(-5, 5, 1000)
cases = []

locs = [0.0]
scales = [1.0]
concentrations = list(jnp.linspace(-2, 2, 9))

for loc, scale, concentration in product(locs, scales, concentrations):
    cases.append({"loc": loc, "scale": scale, "concentration": concentration})


dfs = []
for case in cases:
    gev = tfd.GeneralizedExtremeValue(
        loc=case["loc"], scale=case["scale"], concentration=case["concentration"]
    )

    lp = gev.log_prob(x)
    not_nan_indices = jnp.argwhere(~jnp.isinf(lp))

    nominal_min, nominal_max = nominal_gev_support(
        loc=case["loc"], scale=case["scale"], concentration=case["concentration"]
    )

    data = {
        "nominal_min": nominal_min,
        "nominal_max": nominal_max,
        "observed_min": x[not_nan_indices].min(),
        "observed_max": x[not_nan_indices].max()
    }
    
    data |= case

    df = pd.DataFrame(data, index=[0])
    dfs.append(df)

results = pd.concat(dfs)

results