title: Tensorflow: Accumulating gradients of a Tensor
body: <blockquote>
    <p><strong>TL;DR:</strong> you can just skip to the question in yellow box below.</p>
  </blockquote>
  
  <p>Suppose I have a <strong>Encoder-Decoder</strong> Neural Network, with weights <code>W_1</code> and <code>W_2</code> of the encoder and decoder respectively. Let's denote <code>Z</code> as the output of the encoder. The network is trained with batch size <code>n</code>, and all the gradients will be calculated with respect to the mean loss value over the batch (as shown in image below, the <code>L_hat</code> which is the sum of per-sample loss <code>L</code>).</p>
  
  <p><a href="https://i.stack.imgur.com/BuA7u.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/BuA7u.png" alt="enter image description here"></a></p>
  
  <p>What I'm trying to achieve is, in the backward pass, to <strong>manipulate the gradients</strong> of <code>Z</code> before passing it further to the encoder's weights <code>W_1</code>. Suppose <a href="https://i.stack.imgur.com/Y1pAr.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/Y1pAr.png" alt="&lt;code&gt;mod_grad&lt;/code&gt;"></a> is a somehow modified gradients operator, for which the following holds:</p>
  
  <p><a href="https://i.stack.imgur.com/zo1xr.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/zo1xr.png" alt="enter image description here"></a></p>
  
  <p>The described above, in case of a synchronuous pass (first calculate the modified gradients of <code>Z</code>, then propagate down to <code>W_1</code>) is very easy to implement (the <a href="https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant" rel="nofollow noreferrer">Jacobian</a> multiplication is done using <code>grad_ys</code> of <a href="https://www.tensorflow.org/api_docs/python/tf/gradients" rel="nofollow noreferrer"><code>tf.gradients</code></a>):</p>
  
  <pre class="lang-py prettyprint-override"><code>def modify_grad(grad_z):
      # do some modifications
  
  grad_z = tf.gradients(L_hat, Z)
  mod_grad_z = modify_grad(grad_z)
  mod_grad_w1 = tf.gradients(Z, W_1, mod_grad_z)
  </code></pre>
  
  <blockquote>
    <p>The problem is, <strong>I need to accumulate the gradients</strong> <code>grad_z</code> of the tensor <code>Z</code> over several batches. As the shape of it is dynamic (with <code>None</code> in one of the dimensions, as in the illustration above), I cannot define a <a href="https://www.tensorflow.org/api_docs/python/tf/Variable" rel="nofollow noreferrer"><code>tf.Variable</code></a> to store it. Furthermore, the batch size <code>n</code> may change during training. How can I store the average of <code>grad_z</code> over several batches?</p>
  </blockquote>
  
  <p><strong>PS</strong>: I just wanted to combine pareto-optimal training of <a href="https://arxiv.org/abs/1810.04650" rel="nofollow noreferrer"><code>ArXiv:1810.04650</code></a>, the asynchronous network training of <a href="https://arxiv.org/abs/1609.02132" rel="nofollow noreferrer"><code>ArXiv:1609.02132</code></a>, and batch size scheduling of <a href="https://arxiv.org/abs/1711.00489" rel="nofollow noreferrer"><code>ArXiv:1711.00489</code></a>.</p>