Is there a fast way of calculating 1000 Jacobians of a neural network?

Question

For a numerical implementation I need to calculate the jacobian of a neural network with respect to the input, not the weights, but for 1000 instances in each iteration step that I do. For that I use the tensorflow GradientTape and it looks something like this:

def jacobian(model, unknown, known):
    tJac = time.time()
    with tf.GradientTape() as g:
        x = tf.constant(tf.convert_to_tensor(unknown, dtype=tf.float32))
        h = tf.constant(tf.convert_to_tensor(known, dtype=tf.float32))
        g.watch(x)
        y = tf.convert_to_tensor(model(tf.concat([x, h], 1)), dtype=tf.float32)
    jac = g.jacobian(y, x)
    print('Differentiation took %.3f sec' % (time.time()-tJac))
    return jac

So what I do is that I change the h 1000 times and each time I have to calculate the jacobian and hence have to call the predict of the neural network. This takes 0.16 sec per prediction resulting in 160 seconds since I do it 1000 times. Additionaly since this is only for one iterations step this is not feasible.

What I came up with is that I take my 1000 h instances and put them in a dataframe and then call predict only once. The problem now is that I do not have the ability to use the GradientTape because my code would looke something like this:

def jacobian(input, output):
    tJac = time.time()
    with tf.GradientTape() as g:
        x = tf.constant(tf.convert_to_tensor(input, dtype=tf.float32))
        g.watch(x)
        y = tf.convert_to_tensor(output, dtype=tf.float32)
    jac = g.jacobian(y, x)
    print('Differentiation took %.3f sec' % (time.time()-tJac))
    return jac

This simply results in None.

Answer 1

After trying different things out and reading more about automatic differentation and especially about GradientTape I came across the batch_jacobian function which does exactly what I was looking for.

Might be useful for other people that are also not that familiar with GradientTape.