Finding p% of smallest tensor values - TensorFlow 2.0

Question

I am using Python 3.7 and TensorFlow 2.0 and I have the following problem. Within a given tensor of any dimension, if I want to find p% of the smallest weights (where, p can be a user input), how do I do so?

Example:

x = tf.random.normal(shape=(3, 2, 2))                                   

x.numpy()                                                               
Out[8]: 
array([[[-0.30938825,  0.40093166],
        [ 1.5417175 , -0.9551434 ]],

       [[-0.7977963 , -2.0784302 ],
        [-1.5039488 , -0.75159657]],

       [[ 1.7954558 ,  0.19909047],
        [-0.2359499 , -2.4791834 ]]], dtype=float32)

Within the tensor 'x' (of shape: (3, 2, 2)), how can I find say p = 30% of smallest weights? I then plan on removing the p% of smallest weights by setting them to zero.

Thanks!

Answer 1

Using TensorFlow Probability's stats.percentile:

In [5]: import tensorflow_probability as tfp

In [6]: p30 = tfp.stats.percentile(x, q=30.)

In [7]: p30
Out[7]: <tf.Tensor: id=1159, shape=(), dtype=float32, numpy=-0.9551434>

In [8]: tf.where(x <= p30, 0, x)
Out[8]:
<tf.Tensor: id=1108, shape=(3, 2, 2), dtype=float32, numpy=
array([[[-0.30938825,  0.40093166],
        [ 1.5417175 ,  0.        ]],

       [[-0.7977963 ,  0.        ],
        [ 0.        , -0.75159657]],

       [[ 1.7954558 ,  0.19909047],
        [-0.2359499 ,  0.        ]]], dtype=float32)>

Minimal requirements:

tensorflow==2.0
tensorflow-probability==0.7