Support jit hessian (#397)

* Support jit hessian

* fix

* fix

tests/test_utils.py

@@ -1,4 +1,5 @@
 import unittest
+import torch
 import torchani
 
 
@@ -9,6 +10,9 @@ def testChemicalSymbolsToInts(self):
         self.assertEqual(len(str2i), 6)
         self.assertListEqual(str2i('BACCC').tolist(), [1, 0, 2, 2, 2])
 
+    def testHessianJIT(self):
+        torch.jit.script(torchani.utils.hessian)
+
 
 if __name__ == '__main__':
     unittest.main()

torchani/utils.py

@@ -242,7 +242,13 @@ def __len__(self):
         return len(self.rev_species)
 
 
-def hessian(coordinates, energies=None, forces=None):
+def _get_derivatives_not_none(x: Tensor, y: Tensor, retain_graph: Optional[bool] = None, create_graph: bool = False) -> Tensor:
+    ret = torch.autograd.grad([y.sum()], [x], retain_graph=retain_graph, create_graph=create_graph)[0]
+    assert ret is not None
+    return ret
+
+
+def hessian(coordinates: Tensor, energies: Optional[Tensor] = None, forces: Optional[Tensor] = None) -> Tensor:
     """Compute analytical hessian from the energy graph or force graph.
 
     Arguments:
@@ -263,11 +269,12 @@ def hessian(coordinates, energies=None, forces=None):
     if energies is not None and forces is not None:
         raise ValueError('Energies or forces can not be specified at the same time')
     if forces is None:
-        forces = -torch.autograd.grad(energies.sum(), coordinates, create_graph=True)[0]
+        assert energies is not None
+        forces = -_get_derivatives_not_none(coordinates, energies, create_graph=True)
     flattened_force = forces.flatten(start_dim=1)
     force_components = flattened_force.unbind(dim=1)
     return -torch.stack([
-        torch.autograd.grad(f.sum(), coordinates, retain_graph=True)[0].flatten(start_dim=1)
+        _get_derivatives_not_none(coordinates, f, retain_graph=True).flatten(start_dim=1)
         for f in force_components
     ], dim=1)