Accuracy is lost on laptop's RTX?

[mrkwjc]

Hi! The code below shows how the function jit compiled for a GPU generates significantly different results than the not compiled one or compiled for CPU. This happens on my laptop with RTX A2000, CUDA 12.3, and jax-0.4.24.

import jax

seed = 1234567890
x = jax.random.uniform(jax.random.PRNGKey(seed+10), (2, 2))
w = jax.random.uniform(jax.random.PRNGKey(seed+12), (2, 20))


def func(x, w):
    y = (x + 1.).dot(w)
    return y.sum()


print(func(x, w))
print(jax.jit(func, backend='gpu')(x, w))
print(jax.jit(func, backend='cpu')(x, w))

# 64.80866
# 64.76463
# 64.80866

Is this a bug, or something I should expect? What's interesting, the results are in agreement if one takes x instead of x+1 in the above function...

[mrkwjc]

Problem solved. For better accuracy the higher matmul precision must be set:

jax.config.update('jax_default_matmul_precision', 'high')  # 'bfloat16_3x'