Attempts to produce second derivative vector (diagonal of Hessian) in Theano

theano_second_derivatives.py

import numpy as np
from theano import tensor as T
import theano

x = T.dvector('x')
z = T.dscalar('z')
y = T.sum(T.cos(x*z))

x_test = np.random.normal(size=10)

grad1 = T.grad(y,x)

# Causes an error
grad2, _ = theano.scan(fn=lambda i, x, grad1, z: T.grad(grad1[i], x)[i], 
sequences=T.arange(x.shape[0]), non_sequences=[x, grad1, z])

f = theano.function([x,z], grad2)
print 'Right answer: %s\n computed answer: %s'%(-4*np.cos(x_test*2), f(x_test,2))

Shouldn't this : T.arange(x.shape[0])
be : np.arange(x_test.shape[0])
?