Computing logits for a vector and for all vectors in a set
Why do you want to combine them into two functions? It’s possible, but I don’t think it would be an improvement. They’re two fairly distinct functions, and cramming the code for both into a single function would make the code less readable.

If you really need a single function that handles both, you could do something like:

def softmax(a):
    if isinstance(a, vector):
        return softmax_a_vector(a)
    elif isinstance(a, set):
        return softmax_a_set(a)




which gets what you want, but you keep the nice separation of code into distinct functions.



One other minor thing: you can improve the for loop in softmax_a_set with enumerate:

for idx, a_vector in enumerate(a_set):
    softmax_set[idx] = softmax(a_vector)