使用Tensorflow计算填充张量中的百分位数

如何解决使用Tensorflow计算填充张量中的百分位数

我试图在形状为h.shape = (Batch,Sentences,Words)的张量的最后一个维度中获得大于P = 90％的值，并在Words维度中填充0。我需要的基本上是tf.math.top_k，但又增加了一个困难，即k是一个动态数字，它随Sentences中的每个条目而变化。这意味着k=N/h[:,:,D]对于每个D都是不同的（原因是句子的长度不同）。

h.shape =（批处理，句子，单词）=（4,2,4）的示例输入：

h1 = tf.constant([
[ [0.1,0.7,0.,0.],[0.1,0.] ],[ [0.1,0.3,[0.5,[ [0.3,0.0,0.9,0.4],[0.9,0.] ]])

例如对于[0.3,0.4]，最大数字的P = 50％将是[0.9,0.4]，而对于[0.1,0.]，它将是[0.1]，因为在第一个中我们有4个有效值，并且只有在第二个示例中为1。

现在，执行两个嵌套的while循环，同时为每个句子使用tf.math.count_nonzero确定正确的非0条目数，然后将其重新堆叠为正确的形状！但是，实现的性能速度却要高出10倍以上，所以这是不可行的。

作为参考，如果我们删除张量的前两个维度，此操作将获得我需要的东西：

def maskit(x):
    boolean_mask = tf.cast(x,dtype=tf.bool)
    no_zeros = tf.boolean_mask(x,boolean_mask,axis=0) # problem: flattens x
    pp = tfp.stats.percentile(no_zeros,interpolation='linear',q=90) 
    return pp

maskit(h[0,:]) # correctly returns percentile of [0.1,0.7] instead of [0.1,0.]
maskit(h) # wrongly returns percentile of entire h as if it was flattened

但是输入必须是整个张量h，而输出则必须具有以下形状：(Batch,1)，最后一个暗角保持百分位数。

不幸的是，tf.boolean_mask没有keepdims参数，这就是为什么不能将其应用于整个张量的原因。我还尝试过使用tf.map_fn，tf.slice和类似的运算符。问题总是在于它们采用相等的尺寸，或者永远不保留原始形状。我需要的是一个自定义函数，其功能与tf.math.reduce_mean（但它也不适用于填充张量）或tf.math.count_nonzero和keepdims=True一样运行：在最后一个轴上执行归约运算，编写它，并保持所有其他尺寸不变。

有什么办法可以做到这一点？另外，也许我还缺少一个更笼统的角度：如何在多维填充0的张量上将tf.math.reduce_mean这样的操作与keepdims正确使用？当然，必须有一个简单的解决方案吗？