def _maybe_rotate_dims(self, x, rotate_right=False):
    """"""
    needs_rotation_const = tf.get_static_value(self._needs_rotation)
    if needs_rotation_const is not None and not needs_rotation_const:
      return x
    ndims = prefer_static.rank(x)
    n = (ndims - self._rotate_ndims) if rotate_right else self._rotate_ndims
    perm = prefer_static.concat([
        prefer_static.range(n, ndims), prefer_static.range(0, n)], axis=0)
    return tf.transpose(a=x, perm=perm)