So the v-matrix is fixed?
You could use a nn.Parameter for pre-normalized a, softmax and then broadcasting: v_weighted = v * a.softmax(0).unsqueeze(0)) if you want to weight the columns of v or .unsqueeze(1) if you want to weight the rows. v would be a variable (or a buffer if you want to have it saved in the state_dict).
Best regards
Thomas