A linear (fully connected) layer maps inputs x of shape (..., in_features) to (..., out_features) with weights W of shape (out_features, in_features) and bias b (PyTorch's convention):
Implement linear(x, W, b=None) supporting any number of leading batch dimensions, and linear_flops(batch_shape, in_f, out_f) returning the multiply-adds counted as 2 FLOPs each: 2 * prod(batch_shape) * in_f * out_f.