python - 如何使用 Autograd 在函数内部推导（获取导数）函数？

Question

我有一个看起来像这样的函数：

def f(x, y):
    def g(w):
        def h(z):
            return z * (w ** 2) * (x + y)
        return h
    return g

我想使用autograd计算f wrt (x, y, w, z) 在点 say (x0, y0, w0, z0)的偏导数。 代码当然是python。 谢谢。