Fix documentation build (NVIDIA#8308)

* Fix docs build

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

* Clean up

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

* Fix mock imports

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

* Add comment

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

---------

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

docs/source/conf.py

@@ -64,6 +64,8 @@
     'PIL',
     'boto3',
     'taming',
+    'cytoolz',  # for adapters
+    'megatron',  # for nlp
 ]
 
 _skipped_autodoc_mock_imports = ['wrapt', 'numpy']

docs/update_docs_docker.sh

@@ -1,5 +1,5 @@
 cd ../
-docker run --rm -v $PWD:/workspace python:3.8 /bin/bash -c "cd /workspace && \
+docker run --rm -v $PWD:/workspace python:3.10 /bin/bash -c "cd /workspace && \
 pip install -r requirements/requirements_docs.txt && cd docs/ && rm -rf build && make clean && make html && make html"
 echo "To start web server just run in docs directory:"
 echo "python3 -m http.server 8000 --directory ./build/html/"

nemo/core/neural_types/elements.py

@@ -96,6 +96,10 @@ def fields(self):
         return None
 
     def compare(self, second) -> NeuralTypeComparisonResult:
+        if torch.jit.is_scripting():
+            # Neural types for TorchScript are suppressed
+            # This is a stub to make TorchScript happy
+            return NeuralTypeComparisonResult.SAME
         # First, check general compatibility
         first_t = type(self)
         second_t = type(second)

nemo/core/neural_types/neural_type.py

@@ -51,8 +51,21 @@ def __str__(self):
         else:
             return f"axes: None; elements_type: {self.elements_type.__class__.__name__}"
 
+    def __init__(self, axes: Optional[Any] = None, elements_type: Optional[Any] = None, optional: bool = False):
+        """
+        Args:
+            axes: a tuple of AxisTypes objects representing the semantics of what varying each axis means
+            elements_type: None or ElementType; we need Any annotation here to avoid problems with TorchScript (it is checked in _init_internal)
+            optional: If input to the port of this type can be optional (False by default).
+        """
+        if not torch.jit.is_scripting():
+            self._init_internal(axes=axes, elements_type=elements_type, optional=optional)
+
     @torch.jit.unused
-    def __init__(self, axes: Optional[Any] = None, elements_type: Optional[ElementType] = None, optional=False):
+    def _init_internal(
+        self, axes: Optional[Any] = None, elements_type: Optional[ElementType] = None, optional: bool = False
+    ):
+        """Internals of __init__, separated to make TorchScript and autodoc work"""
         if elements_type is None:
             elements_type = VoidType()
         if not isinstance(elements_type, ElementType):
@@ -62,8 +75,7 @@ def __init__(self, axes: Optional[Any] = None, elements_type: Optional[ElementTy
             )
         self.elements_type = elements_type
         if axes is not None:
-            if not torch.jit.is_scripting():
-                NeuralType.__check_sanity(axes)
+            NeuralType.__check_sanity(axes)
             axes_list = []
             for axis in axes:
                 if isinstance(axis, str):