Stress and PaiNN

Just for future reference, here's a gist that goes through various different ways of computing the stress with autodiff in the dev branch of schnetpack. The expected output (once the straining of the unit cell is fixed), is:

strain
[[[ -2.6313298  -2.980155   -8.266606 ]
  [ -2.980213   -5.880828  -13.080657 ]
  [ -8.266858  -13.08056   -33.894905 ]]]

strain_compute_offsets
[[[ -2.6313148  -2.9801168  -8.266688 ]
  [ -2.9801168  -5.8808103 -13.080613 ]
  [ -8.266622  -13.080922  -33.894863 ]]]

strain_rij
[[[ -2.6312795  -2.9801457  -8.266726 ]
  [ -2.9801457  -5.8808045 -13.080628 ]
  [ -8.266726  -13.080628  -33.894817 ]]]

rij
[[[ -2.6312802  -2.9801457  -8.266726 ]
  [ -2.9801457  -5.880803  -13.0806265]
  [ -8.266726  -13.0806265 -33.89482  ]]]

ase
[[ -2.63129346  -2.98017313  -8.2668025 ]
 [ -2.98017313  -5.88085113 -13.08076411]
 [ -8.2668025  -13.08076411 -33.89520126]]

It's interesting to observe that only the approaches that explicitly use the atom-pair vectors get a perfectly symmetric stress, everything else has some numerical noise.