Practically understanding gradient dimensions on backprop
I've wrote this piece of torch code that implements a Linear(1,1) -&gt; Linear(1,1) -&gt; MSE forward/backward pass:
l1 = torch.nn.Linear(1, 1, bias=False)
l2 = torch.nn.Linear(1, 1, bias=False)
l1.weight = torch.nn.Parameter(
        torch.tensor([[ 0.6279 ]]))
l2.weight = torch.nn.Parameter(
        torch.tensor([[ 0.731 ]]))
loss_fn = torch.nn.MSELoss()
optimizer_l1 = optim.SGD(l1.parameters(), lr=0.001)
optimizer_l2 = optim.SGD(l2.parameters(), lr=0.001)
x = torch.tensor([0.3])
y = torch.tensor([1]).float()
l1_o = l1(x)
l2_o = l2(l1_o)
loss = loss_fn(l2_o, y)
optimizer_l1.zero_grad()
optimizer_l2.zero_grad()
loss.backward()

which outputs that:
l1 weight Parameter containing:
tensor([[0.6279]], requires_grad=True)
l2 weight Parameter containing:
tensor([[0.7310]], requires_grad=True)
======================
x tensor([0.3000])
y tensor([1.])
&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;
l1_o tensor([0.1884], grad_fn=&lt;SqueezeBackward3&gt;)
l2_o tensor([0.1377], grad_fn=&lt;SqueezeBackward3&gt;)
loss tensor(0.7436, grad_fn=&lt;MseLossBackward&gt;)
----------------------
l1_grad tensor([[-0.3782]])
l2_grad tensor([[-0.3249]])


Now, backwards logic of this makes sense:
x = 0.3
l1w = 0.6279
l2w = 0.7310
l1o = 0.1884 -&gt; l2w * x -&gt; 0.6279 * 0.3
l2o = 0.1377

l2_grad = mse_grad(l2o, y, l1o) -&gt; -0.3782
l1_grad = mse_grad(l2o, y, l2w*x) -&gt; -0.3249

However, it stops making sense after we go from 1x1 layers to 2x2:
l1 = torch.nn.Linear(2, 2, bias=False)
l2 = torch.nn.Linear(2, 2, bias=False)
l1.weight = torch.nn.Parameter(
        torch.tensor([[ -0.6279, -0.4686 ], [ -0.0907, 0.6363 ]]))
l2.weight = torch.nn.Parameter(
        torch.tensor([[ 0.731, 0.6026 ], [ -0.1873, -0.6037 ]]))
loss_fn = torch.nn.MSELoss()
optimizer_l1 = optim.SGD(l1.parameters(), lr=0.001)
optimizer_l2 = optim.SGD(l2.parameters(), lr=0.001)
x = torch.tensor([0.3, 0.5])
y = torch.tensor([1, 0]).float()
l1_o = l1(x)
l2_o = l2(l1_o)
loss = loss_fn(l2_o, y)
optimizer_l1.zero_grad()
optimizer_l2.zero_grad()
loss.backward()

Which gives this output:
l1 weight Parameter containing:
tensor([[-0.6279, -0.4686],
        [-0.0907,  0.6363]], requires_grad=True)
l2 weight Parameter containing:
tensor([[ 0.7310,  0.6026],
        [-0.1873, -0.6037]], requires_grad=True)
x tensor([0.3000, 0.5000])
y tensor([1., 0.])
l1_o tensor([-0.4227,  0.2909], grad_fn=&lt;SqueezeBackward3&gt;)
l2_o tensor([-0.1337, -0.0965], grad_fn=&lt;SqueezeBackward3&gt;)
loss tensor(0.6472, grad_fn=&lt;MseLossBackward&gt;)
---------------
l1_grad tensor([[-0.2432, -0.4053],
        [-0.1875, -0.3124]])
l2_grad tensor([[ 0.4792, -0.3298],
        [ 0.0408, -0.0281]])

When I try to reproduce this from scratch, I can't figure out how to keep gradients shape equal to weights shape without messing with my mse_grad function (I want to take gradient with respect to outputs, not weights, this should be possible since they're proportionate):
l1_o = x.dot(l1.T)
l2_o = l1_o.dot(l2.T)
loss = mse(l2_o, y)
l2_o_grad = mse_grad(l2_o, y, l1_o) # shape of l2 (weights) is (2,2), shape of my gradient is (2), since both the layer output and y are of (2) shape

My forward pass, mse loss and mse gradient when called standalone give exact same values as torch.

I'm not a 100% sure if I understood your problem correctly, but the gradient w.r.t a particular layer does have the same dimensionality as that layer's output. The gradient w.r.t (linearly applied) weights in that layer, is not computed by backpropagation but by multiplying the input activations (those before the weights) with the backpropagated gradient. Hence, input activations of [2, 1] multiplied by backpropagated gradients of [1, 2] give you a [2, 2] matrix. Pay attention that the input activations are a vector &quot;standing up&quot; and the backpropagated gradients are &quot;lying on the side&quot;.
I would also recommend you google &quot;The Matrix Cookbook&quot; and familiarize yourself with https://en.wikipedia.org/wiki/Matrix_calculus. Not so much because of the mathematical insights but rather to make sure that derivatives between matrices and vectors and scalars have the right dimensionality (numerator-layout and denominator-layout notation).
