class KMaxPooling(mx.operator.CustomOp):
    def forward(self, is_train, req, in_data, out_data, aux):
        # Desired (k=3):
        # in_data = np.array([1, 2, 4, 10, 5, 3])
        # out_data = [4, 10, 5]
        x = in_data[0].asnumpy()
        idx = x.argsort()[-k:]
        idx.sort(axis=0)
        y = x[idx]