Torch compile + export escn #826
Conversation
Codecov ReportAttention: Patch coverage is
|
| self.use_pbc = use_pbc | ||
| self.use_pbc_single = use_pbc_single | ||
| self.cutoff = cutoff | ||
| self.otf_graph = otf_graph |
There was a problem hiding this comment.
currently this model assumes otf_graph=False. If this is what's going to be in the final version I suggest raising an error if passed otf_graph=True.
There was a problem hiding this comment.
ya let me remove this actually
| energy = torch.zeros(len(natoms), device=node_energy.device) | ||
| energy.index_add_(0, batch_idx, node_energy.view(-1)) | ||
| # Scale energy to help balance numerical precision w.r.t. forces | ||
| energy = energy * 0.001 |
There was a problem hiding this comment.
we want to remove energy = energy * 0.001 eventually but could be kept for backward compat.
|
|
||
| # Compare predicted energies and forces (after inv-rotation). | ||
| energies = out["energy"].detach() | ||
| np.testing.assert_almost_equal(energies[0], energies[1], decimal=5) |
There was a problem hiding this comment.
with reasonable graphs (no max neighbor limit or no strict max neighbor) this should pass with 6 or even 7 decimal places.
| np.testing.assert_array_almost_equal( | ||
| forces[: forces.shape[0] // 2], | ||
| torch.matmul(forces[forces.shape[0] // 2 :], inv_rot), | ||
| decimal=5, |
There was a problem hiding this comment.
with reasonable graphs (no max neighbor limit or no strict max neighbor) this should pass with 6 or even 7 decimal places.
There was a problem hiding this comment.
i tried with energy i can get to 7 decimals but forces only 5
There was a problem hiding this comment.
a single layer should get 7, seems like you are using 8 and that could be why.
| np.testing.assert_array_almost_equal( | ||
| forces[: forces.shape[0] // 2], | ||
| torch.matmul(forces[forces.shape[0] // 2 :], inv_rot), | ||
| decimal=5, |
There was a problem hiding this comment.
a single layer should get 7, seems like you are using 8 and that could be why.
Compilable and Exportable version of escn.
This fully compiles with 0 graph breaks.
Next PR: