From a4f1f1cbc80e7e37f743200af9ddcf4153b2ca28 Mon Sep 17 00:00:00 2001 From: Vladimir Bataev Date: Fri, 2 Feb 2024 01:07:50 +0400 Subject: [PATCH] Fix documentation build (#8308) * Fix docs build Signed-off-by: Vladimir Bataev * Clean up Signed-off-by: Vladimir Bataev * Fix mock imports Signed-off-by: Vladimir Bataev * Add comment Signed-off-by: Vladimir Bataev --------- Signed-off-by: Vladimir Bataev --- docs/source/conf.py | 2 ++ docs/update_docs_docker.sh | 2 +- nemo/core/neural_types/elements.py | 4 ++++ nemo/core/neural_types/neural_type.py | 18 +++++++++++++++--- 4 files changed, 22 insertions(+), 4 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 586f6cf47675..0596b15e3de5 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -64,6 +64,8 @@ 'PIL', 'boto3', 'taming', + 'cytoolz', # for adapters + 'megatron', # for nlp ] _skipped_autodoc_mock_imports = ['wrapt', 'numpy'] diff --git a/docs/update_docs_docker.sh b/docs/update_docs_docker.sh index 19f9b46d3ef1..653894630a3c 100755 --- a/docs/update_docs_docker.sh +++ b/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/" diff --git a/nemo/core/neural_types/elements.py b/nemo/core/neural_types/elements.py index 5d5697d80e46..7e95acebd91f 100644 --- a/nemo/core/neural_types/elements.py +++ b/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) diff --git a/nemo/core/neural_types/neural_type.py b/nemo/core/neural_types/neural_type.py index 49345d6d234f..d00ba72df043 100644 --- a/nemo/core/neural_types/neural_type.py +++ b/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):