Hi! I have RNN outputs h
of shape [T, N]
, and I’m trying to compute the Jacobian of h(t+1)
with respect to the previous state h(t)
(the eigenvalues of such Jacobians are related to chaoticity of the dynamics).
Here’s the code:
def jacobian(h):
"""Compute the Jacobian of an RNN state vector h(t+1) with respect to h(t) for each t.
:param h: shape [T, N]; output from an RNN
:return: tensor of Jacobian matrices of shape [T-1, N, N]
"""
T, N = h.shape
J = []
for t in range(T-1):
J_t = []
for n in range(N):
J_t.append(grad(h[t+1, n], h[t]))
J_t = torch.stack(J_t) # [N, N]
J.append(J_t)
J = torch.stack(J) # [T-1, N, N]
return J
However, running this I get RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
and trying with allow_unused=True
just gives None
s as the gradients…
I guess I don’t really get the autograd that well… is something like this even possible?
EDIT: FYI I don’t need this to be differentiable… just want to look at the eigenvalues “post mortem”.