Chainer で Jacobi 行列を計算

Chainer では多次元ベクトル出力を出す関数のヤコビ行列計算をサポートしていないそうですが、以下のようにすれば計算できますね。

Jacobian.py

def Jacobian(model,x):
    y = model(x)
    n = y.data.size
    for i in range(n):
        x.cleargrad()
        y = model(x)
        z = np.zeros((1,n),dtype=np.float32)
        z[0,i]=1
        y.grad = np.array(z)
        y.backward()
        if i==0:
            J = x.grad
        else:
            J = np.concatenate((J,x.grad))
    return J

in_size = 3
out_size = 2
Wtrue = np.array([[1,1,0],[0,0,2]]).astype(np.float32)
l1 = L.Linear( in_size, out_size, initialW = Wtrue )

x = Variable( np.array([[0,0,0]]).astype(np.float32))

W = Jacobian(l1,x)

出力

>>> array([[ 1.,  1.,  0.],
       [ 0.,  0.,  2.]], dtype=float32)