Difference in precision between vmap'd dot function, and standard version?

[rog77]

Can anyone please explain this behaviour to me? I am still very new to jax, and this is my first attempt with vmap.
Using colab, TPU, I run:

stdfloat = jnp.float32

n = 6
v = 8

vec = device_put(np.random.uniform(low=-1, high=1, size=(n,v)).astype(stdfloat))
mat = device_put(np.random.uniform(low=-1, high=1, size=(n,n)).astype(stdfloat))
output = device_put(np.zeros(shape=(v,n)).astype(stdfloat)) 


for i in range(0, vec.shape[1] ):
  output = output.at[i].set(jnp.dot(mat, vec[:,i],precision=('float32','float32')))


def mydot(mat, vec):
  return jnp.dot(mat, vec, precision=('float32','float32'))

vdot = jax.vmap(mydot, in_axes=(None,1))
Voutput = vdot(mat, vec)

print(output, type(output), output.dtype, "\n")
print(Voutput, type(Voutput), Voutput.dtype)

And, for example, get:

[[-0.84429234  0.86638373 -0.25121966 -1.2597648   0.02286355 -0.1034649 ]
 [ 0.3961082  -0.31665137  0.03001352 -0.29643482 -0.5257444  -0.48592138]
 [-0.68314624  0.62409246  0.31193525 -0.26334375  1.0928401  -0.25894237]
 [-0.6025435  -0.16437384  1.5119045   0.36650294  1.5523392   0.5682912 ]
 [-0.20028746  0.058492    0.25043312 -0.696414    0.750456    0.18312588]
 [-1.251725    0.45785284  0.3905049   0.5865444   0.64460313  0.6542727 ]
 [-0.8486055   0.4602277  -0.43038785  0.07913294 -0.36449912 -0.11145431]
 [-1.0019524   0.02705902 -0.1036967  -0.14137171  0.49729884 -1.0374501 ]] <class 'jax.interpreters.xla._DeviceArray'> float32 

[[-0.84429234  0.8663837  -0.25121963 -1.2597649   0.02286354 -0.10346492]
 [ 0.39610818 -0.31665137  0.03001355 -0.29643482 -0.5257444  -0.48592135]
 [-0.6831463   0.62409246  0.3119353  -0.26334378  1.0928401  -0.2589424 ]
 [-0.6025435  -0.16437383  1.5119045   0.366503    1.5523391   0.5682912 ]
 [-0.2002874   0.05849199  0.25043312 -0.69641405  0.7504559   0.18312584]
 [-1.251725    0.4578528   0.3905049   0.58654445  0.6446031   0.6542726 ]
 [-0.84860545  0.4602277  -0.43038794  0.07913295 -0.3644991  -0.11145433]
 [-1.0019524   0.02705899 -0.10369675 -0.14137171  0.49729887 -1.0374501 ]] <class 'jax.interpreters.xla._DeviceArray'> float32

Is the difference in precision/values my error, or am I missing something with vmap perhaps? I note that when I drop the value of "v" to 1, the discrepancy disappears, e.g.:

[[-0.0789832  -0.2172582  -1.2015206  -1.1114166   0.9628529  -0.21677679]] <class 'jax.interpreters.xla._DeviceArray'> float32 

[[-0.0789832  -0.2172582  -1.2015206  -1.1114166   0.9628529  -0.21677679]] <class 'jax.interpreters.xla._DeviceArray'> float32

Thanks in advance for your help!

[jakevdp]

This is normal floating point precision error: when you add floating point numbers in different orders, the results often have differences of order 1E-8 for single precision arithmetic.

The two operations lower to essentially the same thing (a dot_general primitive call), but in the second case the compiler likely chooses to compute things in a different order. You can reproduce a similar difference without vmap by reversing the entries before computing a matrix-vector product; for example:

import numpy as np
import jax.numpy as jnp
from jax import device_put

stdfloat = jnp.float32

n = 6
v = 8

np.random.seed(0)
vec = device_put(np.random.uniform(low=-1, high=1, size=(n,v)).astype(stdfloat))
mat = device_put(np.random.uniform(low=-1, high=1, size=(n,n)).astype(stdfloat))

def mydot(mat, vec):
  return jnp.dot(mat, vec, precision=('float32','float32'))

out1 = mydot(mat, vec[:, 0])
out2 = mydot(mat[:, ::-1], vec[::-1, 0])

print(out1 - out2)
[ 0.0000000e+00  5.9604645e-08  0.0000000e+00 -5.9604645e-08  0.0000000e+00  2.9802322e-08]

There's not really any way to "fix" this, as it's basically a fact of life when working with floating point arithmetic in any framework.

[rog77]

Thanks for the explanation. While I am thinking about it, I suppose I have to ask: as far as Jax is concerned, if I run the exact same algorithm with the same random key and values (on the same hardware), would I expect whatever the 'error' is to be consistently reproducible, or it it possible that the same numbers might give slightly different answers on another run? I really do appreciate your help :-)

[Jeevesh8]

These errors could compound with rounding:

jnp.round(jnp.array(1.8669497, dtype='float32')*10000)/10000

outputs:

Array(1.8669, dtype=float32)

And

jnp.round(jnp.array(1.8669499, dtype='float32')*10000)/10000

outputs:

Array(1.867, dtype=float32)