Practically understanding gradient dimensions on backprop

I've wrote this piece of torch code that implements a Linear(1,1) -> Linear(1,1) -> 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.])
>>>>>>>>>>>>>>>>>>>>>>
l1_o tensor([0.1884], grad_fn=<SqueezeBackward3>)
l2_o tensor([0.1377], grad_fn=<SqueezeBackward3>)
loss tensor(0.7436, grad_fn=<MseLossBackward>)
----------------------
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 -> l2w * x -> 0.6279 * 0.3
l2o = 0.1377

l2_grad = mse_grad(l2o, y, l1o) -> -0.3782
l1_grad = mse_grad(l2o, y, l2w*x) -> -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=<SqueezeBackward3>)
l2_o tensor([-0.1337, -0.0965], grad_fn=<SqueezeBackward3>)
loss tensor(0.6472, grad_fn=<MseLossBackward>)
---------------
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.