title: How to implement K-Max pooling in Tensorflow or Keras?
body: <p>First off I know that I should use <code>top_k</code> but what makes k-max pooling hard (to implement in TF) is that it has to preserve the order.</p>
  <p>what I have so far:</p>
  <pre><code>import tensorflow as tf
  from tensorflow.contrib.framework import sort
  sess = tf.Session()
  a = tf.convert_to_tensor([[[5, 1, 10, 2], [3, 11, 2, 6]]])
  b = sort(tf.nn.top_k(a, k=2)[1])
  print(tf.gather(a, b, axis=-1).eval(session=sess))
  </code></pre>
  <p>it's close but not there yet</p>
  <p>what I get:</p>
  <pre><code>[[[[[ 5, 10], [ 1, 2]]], [[[ 3, 2], [11, 6]]]]]
  </code></pre>
  <p>what I want:</p>
  <pre><code>[[[5, 10], [11, 6]]]
  </code></pre>
  <p>I am almost hundred percent sure that gather_nd is required but I can't figure that out, also I am a pytorch user and it's really easy there</p>
  <pre><code>import torch
  a = torch.LongTensor([[[5, 1, 10, 2], [3, 11, 2, 6]]])
  b = a.topk(2, dim = -1)[1].sort(dim = -1)[0]
  print(a.gather(-1, b))
  </code></pre>
  <p>Oh and also every code that I found was not an order preserving one(which is semantically wrong)</p>