hoggsumexp

import numpy as np

def hoggsumexp(qns, dqn_dams, diag=False):
    """
    # purpose:
    - Computes L = log(sum(exp(qns, axis=-1))).
    - Also computes its M-dimensional gradient components dL / da_m.

    # input
    - qns: ndarray of shape [n1, n2, n3, ..., nD, N]
    - dqn_dams: ndarray of shape [n1, n2, n3, ..., nD, N, M]
    - diag: if True, then dqn_dams.shape == dqn_dams.shape and [read the source]

    # output
    - L: ndarray of shape [n1, n2, n3, ..., nD]
    - dL_dams: ndarray of shape [n1, n2, n3, ..., nD, M]

    # issues
    - Not exhaustively tested.
    """
    axis = len(qns.shape) - 1
    if diag:
        assert qns.shape == dqn_dams.shape
    Q = np.max(qns)
    expqns = np.exp(qns - Q)
    expL = np.sum(expqns, axis=axis)
    if diag:
        numerator = expqns * dqn_dams
    else:
        numerator = np.sum(np.expand_dims(expqns, axis + 1) * dqn_dams, axis=axis)
    return np.log(expL) + Q, numerator / np.expand_dims(expL, axis)

Calling @dfm : Can you look at this and tell me if I am full of shit?