Multivariate Switch Function

You can use a continuous sigmoid function instead of switch, for example you can see some related discussion here: Switch point Metropolis tuning

Of course, in your case since you need to keep the row sum to 1, it might be a bit more tricky… you might need to construct p in the unconstrained parameter space and transform them row-wise to a simplex