Relationship between (soft)max and (soft)argmax

check_grad_max_softmax.py

import numpy as np


def numeric_grad(f, x, eps=1e-3):

    grad = np.zeros_like(x)

    for i in range(x.shape[0]):
        v = np.zeros_like(x)
        v[i] = 1

        grad[i] = f(x + eps * v) - f(x - eps * v)
        grad[i] /= 2 * eps

    return grad


def logsumexp(x):
    return np.log(np.sum(np.exp(x)))


def softmax(x):
    return np.exp(x) / np.sum(np.exp(x))


def main():

    np.set_printoptions(precision=3, suppress=True)

    x = np.random.randn(5)
    print(f"x: ", x)

    print(f"max(x):", np.max(x))
    print(f"nabla max(x):", numeric_grad(np.max, x))

    print(f"logsumexp(x):", logsumexp(x))
    print(f"nabla logsumexp(x):", numeric_grad(logsumexp, x))
    print(f"softmax(x):", softmax(x))


if __name__ == '__main__':
    main()