有没有可能在Keras中实现一个最大池化层，它提供给定池大小的最大n值？

from keras.layers import Layer
import tensorflow as tf

class MaxNPooling2D(Layer):
    def __init__(self, pool_size, n, **kwargs):
        super(MaxNPooling2D, self).__init__(**kwargs)
        self.pool_size = pool_size
        self.n = n

    def call(self, inputs):
        output = tf.nn.max_pool2d(inputs, ksize=self.pool_size, strides=self.pool_size, padding='VALID')
        output = tf.nn.top_k(output, k=self.n, sorted=True).values
        return output

    def get_config(self):
        config = super(MaxNPooling2D, self).get_config()
        config.update({'pool_size': self.pool_size, 'n': self.n})
        return config

from keras.models import Sequential

model = Sequential()
model.add(MaxNPooling2D(pool_size=(2, 2), n=3))