Is there a layer that takes a tf vector of probabilities and returns a vector with the highest probability is set to 1 and 0 otherwise

Question

I have

layers.Dense(3, activation=\'softmax\')

and this returns a probability for the three classes I