jnp.mean(x, dtype=bfloat16) is not respected

jnp.mean(x, dtype=bfloat16) is not respected: a task in MiMo-V2.6-RL-harbor-code: MiMo-V2.6-RL Code (Harbor) (Harbor dataset). Regarding accumulation dtype for np.mean, the numpy docs say: Note that for floating-point input, the mean is computed using the same precision the input has. Depending on…