Search code examples
pythontensorflowpytorchtensorflow2.0tensorflow-probability

How to get log probabilities in TensorFlow?


I am trying to convert a pytorch script to tensorflow and I need to get log probabilities from a categorical distribution. But the tensorflow calculated log probabilities is different from pytorch's log prob even after using same seed. This is what I have done so far

import torch 
from torch.distributions import Categorical
import tensorflow as tf
import tensorflow_probability as tfp

torch.manual_seed(1)
tf.random.set_seed(1)

probs =[0.4,0.6]
m = Categorical(torch.tensor(probs))
action = m.sample()

n = tfp.distributions.Categorical(probs)
print("pytorch",m.log_prob(action))
print("Tensorflow", tf.math.log(n.prob(action.item())))

Solution

  • tfp.distributions.Categorical(probs)
    

    takes logs as the default argument. They are being normalized and resulting probabilities of built distribution are [.45, .55].

    You need to build tfp distribution as:

     tfp.distributions.Categorical(probs=probs)