contrived_example.py

class Linear:
	def __init__(self, in_dim, out_dim):
		self.params = {
			'W': np.random.rand(in_dim, out_dim)
			'b': np.random.rand(out_dim)
		}

	def run(self, x):
		return sigmoid(self.params['W'] @ x + self.params['b'])


def forward(layer, x):
	y = layer.run(x)

layer = Linear(20, 10)

grad(forward, wrt=layer.params['W'])(layer, np.random.rand(20))