def _get_probs(dist):
    """Helper to extract probability from a distribution."""
    # In tensorflow_probabiliy 0.7 this attribute was filled on __init__ method
    if dist.probs is not None:
        return dist.probs
    # From 0.8 onwards the probs is not filled and a function should be used
    return dist.probs_parameter()