replace np.mean with jnp.mean

replace np.mean with jnp.mean #11