Comment by lowdanie

Since the gradients are computed with jax, the library can be used to differentiate with respect to all inputs including nuclear charge and basis set parameters (exponents and contraction coefficients). I agree that computing gradients for the nuclear charges could be interesting in the context of molecular design.

But for the colab demo I thought that sticking to nuclear positions (i.e atomic forces) would be easier to visualize.