<p>It turns out (per TensorFlow r1.13) that if len(xs) > 1, then tf.hessians(ys, xs) returns tensors corresponding to only the block diagonal submatrices of the full Hessian matrix. Full story and solutions in this paper <a href="https://arxiv.org/pdf/1905.05559" rel="nofollow noreferrer">https://arxiv.org/pdf/1905.05559</a>, and code at <a href="https://github.com/gknilsen/pyhessian" rel="nofollow noreferrer">https://github.com/gknilsen/pyhessian</a></p>
Use tf.gradients or tf.hessians on flattened parameter tensor
<p>Let's say I want to compute the Hessian of a scalar-valued function with respect to some parameters W (e.g the weights and biases of a feed-forward neural network).
If you consider the following code, implementing a two-dimensional linear model trained to minimize a MSE loss: </p>

<pre><code>import numpy as np
import tensorflow as tf

x = tf.placeholder(dtype=tf.float32, shape=[None, 2])  #inputs
t = tf.placeholder(dtype=tf.float32, shape=[None,])  #labels
W = tf.placeholder(np.eye(2), dtype=tf.float32)  #weights

preds = tf.matmul(x, W)  #linear model
loss = tf.reduce_mean(tf.square(preds-t), axis=0) #mse loss

params = tf.trainable_variables() 
hessian = tf.hessians(loss, params)
</code></pre>

<p>you'd expect <code>session.run(tf.hessian,feed_dict={})</code> to return a 2x2 matrix (equal to W). It turns out that because <code>params</code>is a 2x2 tensor, the output is rather a tensor with shape [2, 2, 2, 2]. While I can easily reshape the tensor to obtain the matrix I want, it seems that this operation might be extremely cumbersome when <code>params</code>becomes a list of tensors of varying size (i.e when the model is a deep neural network for instance).</p>

<p>It seems that are two ways around this:</p>

<ul>
<li><p>Flatten <code>params</code> to be a 1D tensor called <code>flat_params</code>:</p>

<pre><code>flat_params = tf.concat([tf.reshape(p, [-1]) for p in params])
</code></pre>

<p>so that <code>tf.hessians(loss, flat_params)</code> naturally returns a 2x2 matrix. However as noted in <a href="https://stackoverflow.com/questions/44836859/why-does-tensorflow-reshape-tf-reshape-break-the-flow-of-gradients">Why does Tensorflow Reshape tf.reshape() break the flow of gradients?</a> for tf.gradients (but also holds for tf.hessians), tensorflow is not able to see the symbolic link in the graph between <code>params</code>and <code>flat_params</code> and <code>tf.hessians(loss, flat_params)</code> will raise an error as the gradients will be seen as <code>None</code>. </p></li>
<li><p>In <a href="https://afqueiruga.github.io/tensorflow/2017/12/28/hessian-mnist.html" rel="nofollow noreferrer">https://afqueiruga.github.io/tensorflow/2017/12/28/hessian-mnist.html</a>, the author of the code goes the other way, and first create the flat parameter and reshapes its parts into <code>self.params</code>. This trick does work and gets you the hessian with its <em>expected</em> shape (2x2 matrix). However, it seems to me that this will be cumbersome to use when you have a complex model, and impossible to apply if you create your model via built-in functions (like <code>tf.layers.dense</code>, ..).  </p></li>
</ul>

<p>Is there no straight-forward way to get the Hessian matrix (as in the 2x2 matrix in this example) from <code>tf.hessians</code>, when <code>self.params</code> is a list of tensor of arbitrary shapes? If not, how can you automatize the reshaping of the output tensor of <code>tf.hessians</code>?</p>