Rate this Page
★ ★ ★ ★ ★

torch.utils.flop_counter.sdpa_flop_count#

torch.utils.flop_counter.sdpa_flop_count(query_shape, key_shape, value_shape)[source]#

Count flops for self-attention.

NB: We can assume that value_shape == key_shape

On this page
Show Source
PyTorch Libraries