Как получить тот же эффект и в чем разница между тензорным потоком.multinomial и pytorch.multinomial?

Question

tensorflow.multinomial и pytorch.multinomial, разрыв в производительности огромен. В частности, я думаю, что tennsflow.multinomial очень проблематичен.

Я написал следующий код для сравнения двух функций, я установил для первого числа 10, а для остальных - 0,0001, а затем Softmax в качестве вероятности выбора.

import torch
import numpy as np
import  tensorflow as tf

single_data = np.array([10,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001,0.0001])


single_data = np.expand_dims(single_data, axis=0)
data =np.repeat(single_data, 50, axis=0)
data = tf.nn.softmax(data,axis=-1)
sampled_word = tf.multinomial(data,1)
sample = tf.reshape(sampled_word,[-1])

with tf.Session() as sess:
    a_data=sess.run(data)
    print("prob:",a_data)
    print("tensorflow.multinomial", sess.run(sample))

    a_data = torch.from_numpy(a_data)
    idx = torch.multinomial(a_data, num_samples=1)
    print("tensorflow.multinomial", np.reshape(idx.data.numpy(),[-1]))

prob: [[9.93907740e-01 4.51278541e-05 4.51278541e-05 ... 4.51278541e-05
  4.51278541e-05 4.51278541e-05]
 [9.93907740e-01 4.51278541e-05 4.51278541e-05 ... 4.51278541e-05
  4.51278541e-05 4.51278541e-05]
 [9.93907740e-01 4.51278541e-05 4.51278541e-05 ... 4.51278541e-05
  4.51278541e-05 4.51278541e-05]
 ...
 [9.93907740e-01 4.51278541e-05 4.51278541e-05 ... 4.51278541e-05
  4.51278541e-05 4.51278541e-05]
 [9.93907740e-01 4.51278541e-05 4.51278541e-05 ... 4.51278541e-05
  4.51278541e-05 4.51278541e-05]
 [9.93907740e-01 4.51278541e-05 4.51278541e-05 ... 4.51278541e-05
  4.51278541e-05 4.51278541e-05]]
tensorflow.multinomial [121 129  35 104   4 133  60  92 104 129   4  49  35  99 109 111  62  87
  23   5 109  63 103  61  78  43 101  85   2 127   0  36  53   0  74  44
  64  55  51  59 108   0 112  32  36  24  68 135  72  22]
tensorflow.multinomial [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0]

Мы обнаружили, что Tensorflow было почти трудно выбирать данные (99% времени), в то время как PyTorch легко выбирался. Любой комментарий?