class MissingGradients(Module):
def __init__(self) -> None:
super().__init__()
self.fc1 = Linear(784, 30)
self.fc2 = Linear(30, 10)
self.set_parameters("fc1", "fc2")
def forward(self, x: NPArray) -> NPArray:
x = self.fc1(x)
x = self.fc2(x)
return x
def backward(self, dout: NPArray) -> dict[str, NPArray]:
dx2, _, _ = self.fc2.backward(dout)
_, _, _ = self.fc1.backward(dx2)
return {}
class MissingParameters(Module):
def __init__(self) -> None:
super().__init__()
self.fc1 = Linear(784, 30)
self.fc2 = Linear(30, 10)
def forward(self, x: NPArray) -> NPArray:
x = self.fc1(x)
x = self.fc2(x)
return x
def backward(self, dout: NPArray) -> dict[str, dict[str, NPArray]]:
grads = {}
dx2, dw2, db2 = self.fc2.backward(dout)
grads["fc2"] = {"w": dw2, "b": db2}
_, dw1, db1 = self.fc1.backward(dx2)
grads["fc1"] = {"w": dw1, "b": db1}
return grads
def test_linear_backward() -> None:
x = np.random.randn(10, 2, 3)
w = np.random.randn(6, 5)
b = np.random.randn(5)
dout = np.random.randn(10, 5)
model = Linear(10, 5)
params = {"w": w, "b": b}
dx_num, dw_num, db_num = estimate_gradients(model, dout, x, params)
_ = model(x)
dx, dw, db = model.backward(dout)
assert np.allclose(dx_num, dx)
assert np.allclose(dw_num, dw)
assert np.allclose(db_num, db)