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.