tf.random.categorical giving strange results

I am trying to implement np.random.choice in tensorflow. Here is my implementation

import numpy as np
import tensorflow as tf
p=tf.Variable(0,tf.int32)
selection_sample=[i for i in range(10)]#sample to select from
k=tf.convert_to_tensor(selection_sample)
samples = tf.random.categorical(tf.math.log([[1, 0.5, 0.3, 0.6]])
sample_selected=tf.cast(samples[0][0],tf.int64)
op=tf.assign(p,k[sample_selected])
#selection_sample[samples]
init=tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init)
    print(sample_selected.eval())
    print(k.eval())
    print((sess.run(op)))
    print(p.eval())

However when sample_selected is for example 1 i expect p.eval to be 1 i.e k[1] but this is not the case. For example running this code a sample output is

3
[0 1 2 3 4 5 6 7 8 9]
1
1

yet p.eval should be k[3] and sess.run(op) should also be k[3] what I am doing wrong. Thanks