Display metatensor-torch torchscript classes correctly in docs (#40)

* Display torchscript classes correctly in docs
* Disable compiled functions when building docs

docs/src/conf.py

@@ -10,6 +10,10 @@
 ROOT = os.path.abspath(os.path.join("..", ".."))
 sys.path.insert(0, ROOT)
 
+# when importing metatensor-torch, this will change the definition of the classes
+# to include the documentation
+os.environ["METATENSOR_IMPORT_FOR_SPHINX"] = "1"
+
 # -- Project information -----------------------------------------------------
 
 # The master toctree document.

src/metatensor/models/utils/data/dataset.py

@@ -1,3 +1,4 @@
+import os
 from typing import Dict, List
 
 import metatensor.torch
@@ -6,8 +7,13 @@
 from metatensor.torch.atomistic import ModelCapabilities, System
 
 
-compiled_slice = torch.jit.script(metatensor.torch.slice)
-compiled_join = torch.jit.script(metatensor.torch.join)
+if os.environ.get("METATENSOR_IMPORT_FOR_SPHINX", "0") == "1":
+    # This is necessary to make the Sphinx documentation build
+    compiled_slice = None
+    compiled_join = None
+else:
+    compiled_slice = torch.jit.script(metatensor.torch.slice)
+    compiled_join = torch.jit.script(metatensor.torch.join)
 
 
 class Dataset(torch.utils.data.Dataset):