import math


def calculate_gain(nonlinearity, param=None):
    """Return the gain value according to nonlinearity function.

    Parameters
    ----------
    nonlinearity : str
        The nonlinearity to compute gain value.
    param : number, optional
        The optional param value for nonlinearity.

    Returns
    -------
    number
        The gain value.

    """
    linear_fns = [
        'linear',
        'conv1d',
        'conv2d',
        'conv3d',
        'conv_transpose1d',
        'conv_transpose2d',
        'conv_transpose3d',
    ]
    if nonlinearity in linear_fns or nonlinearity == 'sigmoid':
        return 1
    elif nonlinearity == 'tanh':
        return 5.0 / 3
    elif nonlinearity == 'relu':
        return math.sqrt(2.0)
    elif nonlinearity == 'leaky_relu':
        if param is None:
            negative_slope = 0.01
        else:
            try:
                negative_slope = float(param)
            except ValueError:
                raise ValueError("Negative slope {} is not a valid number".format(param))
        return math.sqrt(2.0 / (1 + negative_slope ** 2))
    else:
        raise ValueError('Unsupported nonlinearity: ' + nonlinearity)