title: How to get loss gradient wrt internal layer output in tensorflow 2?
body: <p>I would like to get gradient of the model's loss function with respect to specific layer's output during training. What I want to do with it next, is using a value of that gradient to modify something in layer in the next learning epoch. 
  So how to obtain that gradient?</p>
  
  <p>Here's a minimal example. 
  MinimalRNNCell code is copied from TensorFlow's website and toy data is provided only to reproduce the behavior.</p>
  
  <pre><code>import tensorflow as tf 
  from tensorflow.keras.layers import RNN, SimpleRNNCell, SimpleRNN, Layer, Dense, AbstractRNNCell
  from tensorflow.keras import Model
  import numpy as np
  import tensorflow.keras.backend as K
  
  
  class MinimalRNNCell(AbstractRNNCell):
  
      def __init__(self, units, **kwargs):
        self.units = units
        super(MinimalRNNCell, self).__init__(**kwargs)
  
      @property
      def state_size(self):
        return self.units
  
      def build(self, input_shape):
        self.kernel = self.add_weight(shape=(input_shape[-1], self.units),
                                      initializer='uniform',
                                      name='kernel')
        self.recurrent_kernel = self.add_weight(
            shape=(self.units, self.units),
            initializer='uniform',
            name='recurrent_kernel')
        self.built = True
  
      def call(self, inputs, states):
        prev_output = states[0]
        h = K.dot(inputs, self.kernel)
        output = h + K.dot(prev_output, self.recurrent_kernel)
        return output, output
  
  
  class MyModel(Model):
      def __init__(self, size):
          super(MyModel, self).__init__()
          self.minimalrnn=RNN(MinimalRNNCell(size), name='minimalrnn')
          self.out=Dense(4)
  
      def call(self, inputs):
          out=self.minimalrnn(inputs)
          out=self.out(out)
          return out
  
  
  x=np.array([[[3.],[0.],[1.],[2.],[3.]],[[3.],[0.],[1.],[2.],[3.]]])
  y=np.array([[[0.],[1.],[2.],[3.]],[[0.],[1.],[2.],[3.]]])
  
  model=MyModel(2)
  model.compile(optimizer='sgd', loss='mse')
  model.fit(x,y,epochs=10, batch_size=1, validation_split=0.2)
  
  
  
  </code></pre>
  
  <p>Now I want to get gradient of output of MyModel's minimalrnn layer (after every batch of data).</p>
  
  <p>How to do this? I suppose I can try with GradientTape watching model.get_layer('minimalrnn').output, but I need more learning resources or examples. </p>
  
  <p><strong>EDIT</strong></p>
  
  <p>I used GradientTape as in code provided by Tiago Martins Peres, but I specifically want to obtain gradient wrt layer output, and I'm still not able to achieve that. </p>
  
  <p>Now after class definitions my code looks like this:</p>
  
  <pre><code>
  x=np.array([[[3.],[0.],[1.],[2.],[3.]],[[3.],[0.],[1.],[2.],[3.]]])
  y=np.array([[0., 1., 2., 3.],[0., 1., 2., 3.]])
  
  model=MyModel(2)
  
  #inputs = tf.keras.Input(shape=(2,5,1))
  #model.call(x)
  
  def gradients(model, inputs, targets):
      with tf.GradientTape() as tape:
          tape.watch(model.get_layer('minimalrnn').output)
          loss_value = loss_fn(model, inputs, targets)
      return tape.gradient(loss_value, model.trainable_variables)
  
  def loss_fn(model, inputs, targets):
      error = model(inputs) - targets
      return tf.reduce_mean(tf.square(error))
  
  optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
  print("Initial loss: {:.3f}".format(loss_fn(model, x, y)))
  for i in range(10):
      grads = gradients(model, x, y)
      optimizer.apply_gradients(zip(grads, model.trainable_variables))
      print("Loss at step {:03d}: {:.3f}".format(i, loss_fn(model, x, y)))
  print("Final loss: {:.3f}".format(loss_fn(model, x, y)))
  </code></pre>
  
  <p>As you can see I added tape.watch in gradients function definition, because I want to watch layer output. However I'm getting error:</p>
  
  <pre><code>Traceback (most recent call last):
    File "/home/.../test2.py", line 73, in &lt;module&gt;
      grads = gradients(model, x, y)
    File "/home/.../test2.py", line 58, in gradients
      print(model.get_layer('minimalrnn').output)
    File "/home/.../.venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 1553, in output
      raise AttributeError('Layer ' + self.name + ' has no inbound nodes.')
  AttributeError: Layer minimalrnn has no inbound nodes.
  </code></pre>
  
  <p>I also tried to call model on Input with specified size (commented lines), according to answer to this: <a href="https://stackoverflow.com/questions/59383356/accessing-layers-input-output-using-tensorflow-2-0-model-sub-classing">Accessing layer&#39;s input/output using Tensorflow 2.0 Model Sub-classing</a>. It didn't help. Specifying input shape in model's init function, like below,  also doesn't help - still the same error.</p>
  
  <pre><code>self.minimalrnn=RNN(MinimalRNNCell(size), name='minimalrnn', input_shape=(2,5,1))
  </code></pre>