keras/src/layers/attention/attention.py/Attention/compute_output_shape
class Attention:    def compute_output_shape(self, input_shape):
        """Returns shape of value tensor dim, but for query tensor length"""
        return (*input_shape[0][:-1], input_shape[1][-1])