title: Keras custom layer using argmax kills gradients
body: <p>After implementing a custom layer in Keras and using it in my network, I realized that all previous layers are not updated during training anymore. In fact, analyzing the <em>tensorboard distributions</em> of the weights and gradients computed for my network, the gradients for all layers preceding the custom layer are zero. </p>
  
  <p>The layer I am trying to implement performs the following:</p>
  
  <p><strong>Goal:</strong>
  Given a tensor of size <code>B x H x W x D</code>, with <code>D=1</code>, add two new channels of size <code>H x W</code> to each sample of the batch (the values of these channels are given as a constant input, so they are not learned or dynamic). For the resulting tensor of shape <code>B x H x W x 3</code>, compute the <em>argmax</em> on the first channel sample-wise (i.e. over the 2 dimensional image) and apply the resulting indices on the second and third channel.</p>
  
  <p><strong>Shapes:</strong>
  Given <code>B x H x W x D = (B x 8 x 8 x 1)</code>, the layer outputs a tensor of shape <code>B x 2 x 1</code> with an intermediate representation of size <code>B x H x W x 3</code>.</p>
  
  <p><strong>Numerical example:</strong>
  Let's assume <code>B = 1</code>, <code>H = W = 2</code>, and let the input tensor have the following values:</p>
  
  <pre><code>input[0, :, :, 0] = [[0.1, 0.2],
                       [0.3, 0.4]]
  </code></pre>
  
  <p>Adding the new channels:</p>
  
  <pre><code>input[0, :, :, 0] = [[0.1, 0.2],
                       [0.3, 0.4]]
  
  input[0, :, :, 1] = [[a, b],
                       [c, d]]
  
  input[0, :, :, 2] = [[e, f],
                       [g, h]]
  </code></pre>
  
  <p>The filter computes the <em>argmax</em> on the first channel, and returns the corresponding elements of the second and third channel:</p>
  
  <pre><code>indices = argmax(input[0, :, :, 0]) = [1, 1]
  input[0,indices[0], indices[1], 1] = [d]
  input[0,indices[0], indices[1], 1] = [h]
  
  Layer(input) = [
                   [[d], [h]]
                 ]
  </code></pre>
  
  <p>Here I attach the code I wrote to implement the filter's logic:</p>
  
  <pre><code>class CustomGlobalMaxPooling(Layer):
  
      def __init__(self, **kwargs):
          """Layer computing the argmax in the first channel of each sample in the input
          tensor and returning the values associated to the second and third channel.
          """
  
          super(CustomGlobalMaxPooling, self).__init__(**kwargs)
  
      def build(self, input_shape):
          assert len(input_shape) == 4, "Only 4 dimensional input supported (B, H, W, D)"
          assert input_shape[-1] == 1, "Only supports one channels. Given shape: {}".format(input_shape)
          super(CustomGlobalMaxPooling, self).build(input_shape)
  
      def call(self, x, **kwargs):
          batch_size = tf.shape(x)[0]
          r = tf.cast(tf.range(batch_size), tf.int64)
  
          img_h = tf.shape(x)[1]
          img_w = tf.shape(x)[2]
  
          ## HERE WE ADD THE NEW CHANNELS ON EACH SAMPLE (code omitted) ##
          x = add_channels(x)
  
          flattened_values = tf.reshape(x[:, :, :, 0], (batch_size, img_h * img_w))
          flattened_1 = tf.reshape(x[:, :, :, 1], (batch_size, img_h * img_w))
          flattened_2 = tf.reshape(x[:, :, :, 2], (batch_size, img_h * img_w))
  
          argmax = tf.argmax(flattened_values, axis=1)
          argmax = tf.transpose(tf.stack([r, argmax]), [1, 0])
  
          max_1 = tf.gather_nd(flattened_1, argmax)
          max_2 = tf.gather_nd(flattened_2, argmax)
  
          max_1 = tf.expand_dims(max_1, -1)
          max_2 = tf.expand_dims(max_2, -1)
  
          return tf.transpose(tf.stack((max_1, max_2)), [1, 0, 2])
  
      def compute_output_shape(self, input_shape):
          batch_size = input_shape[0]
          return (batch_size, 2, 1)
  </code></pre>
  
  <p>I guess that by performing the armgax on one channel and applying the indices on new channels, my network is not differentiable anymore, so no gradients can propagate.</p>
  
  <p>Is there a way around this problem?</p>