求助：TensorFlow中仅计算指定位置Jacobian矩阵的优化方法

我正在为用于自回归的神经网络计算Jacobian矩阵。该模型有9个输入变量，输出3个变量，输入shape=(1,9)，输出shape=(1,3)。当前通过TensorFlow代码可计算出维度为(3,9)的完整Jacobian矩阵，但计算速度较慢，且我仅需计算矩阵中示意图标记位置的元素。此Jacobian矩阵用于扩展卡尔曼滤波，现附上相关代码及示意图说明需求：

为清晰说明需求，附上需计算的Jacobian位置示意图：图中o/p代表输出变量，i/p代表输入变量，数字为变量位置，层中数字为隐藏层神经元数量。

Jacobian计算代码 #

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)

神经网络代码 #

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))

内容的提问来源于stack exchange，提问作者rahul