[Numerics] semi-local functionals from `jax-xc` only works with `float64`

[zekun-shi]

When using semi-local functionals (e.g. gga_x_pbe) under float32, on some grid point we will have NaN gradients w.r.t. to mo_coeff. If float64 is used we won't get this problem. Some numerical tuning is required.