modern_hopfield.py

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


class ModernHopfield:
    def __init__(self, memories):
        self.memories = memories

    def update(self, query):
        normalised_memories = self.memories / np.vstack(np.linalg.norm(self.memories, 
                                                                       axis=1))
        return softmax(normalised_memories @ query) @ self.memories