el_ion_diff – shape [N_batch x n_el x N_ion x 3]
el_ion_dist – shape [N_batch x n_el x N_ion]
orbitals –
Returns: