Updates MIMO's DenseMultiHead layer to support TPUs.
Updates MIMO's DenseMultiHead layer to support TPUs.

Previously, runs would fail on TPUs with the error message `Dynamic input
dimension to reshape that is both splitted and combined is not supported`.

PiperOrigin-RevId: 413757034

google/edward2