Jacobian of tensorflow model

I am trying to calculate jacobian matrix from my neural network trained for autoregression.There are 9 input variables to the model and it predicts 3 variables as output.
Input shape=(1,9)
Output shape=(1,3)
And I can calculate the normal jacobian matrix of shape (3,9) from the current code of tensorflow.

I have the representation for the matrix that I can currently calculate attached here Jacobian Matrix.

My issue is this jacobian calculation is very slow and I don't want to calculate all the jacobians and only those jacobians at the place marked in the above image.

I have the code snippets relevant to this issue.This jacobian for is used for extended kalman filter. Can some one help me figure out how can I do this in tensorflow.

Code for jacobian calculation

 def jacobian_tensorflow(self,verbose=False): 
        jacobian_matrix = []
        it = tqdm(range(self.output_size)) if verbose else range(self.output_size)
        for o in it:
            grad_func = tf.gradients(self.nn_model.output[:, o], self.nn_model.input)
            gradients = sess.run(grad_func, feed_dict={self.nn_model.input: self.pred_x.reshape((1, self.pred_x.size))})
            jacobian_matrix.append(gradients[0][0,:])
        return np.array(jacobian_matrix)

Code of my neural networks

input_window = Input(shape=(deg_order * 3,))
    x = Dense(90, activation='tanh')(input_window)
    x = Dense(60, activation='tanh')(x)
    x = Dense(30, activation='tanh')(x)
    x = Dense(15, activation='tanh')(x)
    output = Dense(3, activation='tanh')(x)
    autoencoder_model = Model(input_window, output)

    autoencoder_model.compile(optimizer='adam', loss=tf.keras.metrics.mean_squared_error)

    autoencoder_model.fit(x_train, y_train, epochs=epochs,
                          shuffle=True,
                          validation_data=(x_validate, y_validate))

To explain clearly I have added this image of what jacobian i want to calculate.

In image o/p means the output variable and i/p is the variable on the input side. The number are just for position of the variables on the input and output side. The numbers in layers are the neurons in that hidden layer Jacobian I want to calculate from the neural network