keras/src/layers/attention/multi_head_attention.py/_get_output_shape
def _get_output_shape(output_rank, known_last_dims):
    return [None] * (output_rank - len(known_last_dims)) + list(known_last_dims)