1st sample from tf.multinomial is always zero in webgpu backend

[alvinleung1996]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow.js): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS 13.6.1 (22G313)
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A
• TensorFlow.js installed from (npm or script link): script link
• TensorFlow.js version (use command below): 4.12.0
• Browser version: Version 119.0.6045.105 (Official Build) (x86_64)
• Tensorflow.js Converter Version: N/A

Describe the current behavior

When sampling from tf.multinomial using the webgpu backend, the 1st sample is always zero.

Describe the expected behavior

All samples should respect the the given logits.

Standalone code to reproduce the issue

Execute the following code in the console in https://js.tensorflow.org/api/4.12.0/

await tf.setBackend('webgpu')

let sum = tf.zeros([10])

const logits = tf.tensor1d([0.1, 0.9]).log()
for (let i = 0; i < 1000; i++) {
  const samples = tf.multinomial(logits, 10)
  sum = sum.add(samples)
}

sum.print()

Other info / logs

webgpu backend:

Tensor
    [0, 894, 886, 900, 900, 912, 887, 877, 896, 880]

webgl backend:

Tensor
    [895, 901, 896, 884, 914, 894, 908, 906, 907, 912]

[gaikwadrahul8]

Hi, @alvinleung1996

Thank you for bringing this issue to our attention and I tried to replicate the same issue from my end and I'm also getting same behavoiur which you mentioned in the issue template so we'll have to dig more into this issue and will update you soon. Thank you!

For reference I have added screenshot below:

WebGPU Backend output :

WebGL Backend output :