That is why (in the diagonal case), you can create a parameter w with only the size:
[out_channels, in_channels, kernel_size] (for example : [16, 32, 3]).
This parameter will only contain the useful diagonal information, without the useless zeros. So at the optimization step, only these floats will be updated.
It is just the tensor torch.diag_embed(w) which contains also the zeros. This one has the size:
[out_channels, in_channels, kernel_size, kernel_size] (for example : [16, 32, 3, 3]).
And in the general case, I’m sure you can always find a way to create a parameter with only the useful floats (that need to be updated), and then transform it to the form you need.