jacobian_of_solve.py

import theano
import numpy as np

# inti vars
s = theano.tensor.matrix('s')
y = theano.tensor.vector('y')
r = theano.tensor.slinalg.Solve()(s, y)

# jacobian using scan
dr1 = theano.tensor.jacobian(r.flatten(), s)
# jacobian using Lop and identity matrix
dr2 = theano.tensor.Lop(r, s, theano.tensor.eye(s.shape[0]))

# compile funcs
f = theano.function(outputs=r, inputs=[s, y])
df1 = theano.function(outputs=dr1, inputs=[s, y])
df2 = theano.function(outputs=dr2, inputs=[s, y])

# test
D = 200
S = np.random.randn(D, D).astype(theano.config.floatX)
S = S + S.T
Y = np.ones(D).astype(theano.config.floatX)
R1 = df1(S, Y)
R2 = df2(S, Y)
np.allclose(R1, R2.reshape(R1.shape), atol=1e-4, rtol=1e-4)