how to pmap sync batch_stats? maybe a bug?

[zhenlan0426]

@cgarciae kindly point me to this example for pmap with batchNorm. And the batch_stats are sync across devices with pmean.

  # sync batch stats
  batch_stats = jax.lax.pmean(variables['batch_stats'], "device")
  variables = variables.copy({'batch_stats': batch_stats})

And this is how batch_stats is updated inside BatchNorm module,

ra_mean.value = self.momentum * ra_mean.value + (1 -self.momentum) * mean
ra_var.value = self.momentum * ra_var.value + (1 - self.momentum) * var

The pmean of var across device wont be the same as the true var calculated over the entire data (device * batch) if we dont divided the data into device. This way the variance stats in batchNorm seems off.

[zhenlan0426]

This simple example shows this empirically.

import numpy as np
x = np.random.randn(128)
# var over the entire data, var over device first and then average across device
print(np.var(x),x.reshape(-1,8).var(0).mean())