def probabilities_from_activation(self, vmap): return activation_functions.softmax(vmap[self])
def sample_from_activation(self, vmap): p = activation_functions.softmax(vmap[self]) return samplers.multinomial(p)