|
21 | 21 |
|
22 | 22 | import jaxtyping |
23 | 23 | import sphinx_rtd_theme # noqa |
24 | | -from uncompyle6.semantics.fragments import code_deparse |
25 | 24 |
|
26 | 25 |
|
27 | 26 | sys.path.append(os.path.abspath(os.path.join(__file__, "..", "..", ".."))) |
|
63 | 62 | "tasklist", # Check boxes |
64 | 63 | ] |
65 | 64 |
|
| 65 | +suppress_warnings = [ |
| 66 | + # We use subsections of README in docs, which start with a lower header level |
| 67 | + # than H1, which makes myst_parser complain. This suppresses these warnings. |
| 68 | + "myst.header", |
| 69 | + # The config includes the _process function below, which is "unpickable". |
| 70 | + # This suppresses warnings about caching such values. |
| 71 | + "config.cache", |
| 72 | +] |
| 73 | + |
66 | 74 | # Add any paths that contain templates here, relative to this directory. |
67 | 75 | templates_path = ["_templates"] |
68 | 76 |
|
|
87 | 95 | # Example configuration for intersphinx: refer to the Python standard library. |
88 | 96 | intersphinx_mapping = { |
89 | 97 | "python": ("https://docs.python.org/3/", None), |
90 | | - "torch": ("https://pytorch.org/docs/master/", None), |
| 98 | + "torch": ("https://pytorch.org/docs/main/", None), |
91 | 99 | } |
92 | 100 |
|
93 | 101 | # Disable documentation inheritance so as to avoid inheriting |
@@ -117,15 +125,15 @@ def _convert_internal_and_external_class_to_strings(annotation): |
117 | 125 |
|
118 | 126 | # Convert jaxtyping dimensions into strings |
119 | 127 | def _dim_to_str(dim): |
120 | | - if isinstance(dim, jaxtyping.array_types._NamedVariadicDim): |
| 128 | + if isinstance(dim, jaxtyping._array_types._NamedVariadicDim): |
121 | 129 | return "..." |
122 | | - elif isinstance(dim, jaxtyping.array_types._FixedDim): |
| 130 | + elif isinstance(dim, jaxtyping._array_types._FixedDim): |
123 | 131 | res = str(dim.size) |
124 | 132 | if dim.broadcastable: |
125 | 133 | res = "#" + res |
126 | 134 | return res |
127 | | - elif isinstance(dim, jaxtyping.array_types._SymbolicDim): |
128 | | - expr = code_deparse(dim.expr).text.strip().split("return ")[1] |
| 135 | + elif isinstance(dim, jaxtyping._array_types._SymbolicDim): |
| 136 | + expr = dim.elem |
129 | 137 | return f"({expr})" |
130 | 138 | elif "jaxtyping" not in str(dim.__class__): # Probably the case that we have an ellipsis |
131 | 139 | return "..." |
@@ -220,4 +228,4 @@ def _process(annotation, config): |
220 | 228 | # @jpchen's hack to get rtd builder to install latest pytorch |
221 | 229 | # See similar line in the install section of .travis.yml |
222 | 230 | if "READTHEDOCS" in os.environ: |
223 | | - os.system("pip install torch==1.11.0+cpu -f https://download.pytorch.org/whl/torch_stable.html") |
| 231 | + os.system("pip install torch==2.0+cpu -f https://download.pytorch.org/whl/torch_stable.html") |
0 commit comments