如何使用张量流优化反演采样？

Question

感谢您的帮助。

我正在进行CDF反转采样。

算法如下所示。

---

ㆍ输入：A = [78，226，344，425，473，495，506，510，511]

ㆍ选择随机整数矩阵rnd [i] [j]。 // 0 <= rnd <512，i = 536，j = 256

ㆍ找到最小整数S = A [k]的索引，使得rnd [i] [j] <= A [k] // 0 <= k

ㆍ输出：S [i] [j] // uint16类型

我在下面的代码中使用tensorflow。

rnd=tf.random_uniform([536, 256], minval=0, maxval=512, dtype=tf.int32, seed=None, name=None) 

cdf1=tf.fill([536, 256], 8)
r1=tf.where(tf.less(rnd, 511), cdf1-1, cdf1)
r2=tf.where(tf.less(rnd, 507), cdf1-2, r1)
r3=tf.where(tf.less(rnd, 496), cdf1-3, r2)
r4=tf.where(tf.less(rnd, 474), cdf1-4, r3)
r5=tf.where(tf.less(rnd, 426), cdf1-5, r4)
r6=tf.where(tf.less(rnd, 345), cdf1-6, r5)
r7=tf.where(tf.less(rnd, 227), cdf1-7, r6)
r8=tf.where(tf.less(rnd, 79), cdf1-8, r7)

S=tf.cast(r8, tf.uint16)

但这太慢了。

如何快速计算S？