def normalize(a, axis=None):
    """Normalizes the input array so that it sums to 1.

    :param a: non-normalized input array
    :param axis: dimension along which normalization is to be performed
    """
    a_sum = a.sum(axis)
    if axis and a.ndim > 1:
        a_sum[a_sum == 0] = 1
        shape = list(a.shape)
        shape[axis] = 1
        a_sum.shape = shape
    a /= a_sum
