Add docs on implementing PyTorch Ops (and CumOp) (#837)

* Add Pytorch support for Cum Op * Modify test for Cum op * Raise TypeError if axis not int or None * Fix init method of CumOp * Extend tutorial on documentation for Pytorch * Add tab for Numba * Add intersphinx mapping * Parametrize dtype
pymc-devs · Jul 4, 2024 · 781073b · 781073b
1 parent e57e25b
commit 781073b
Show file tree

Hide file tree

Showing 7 changed files with 550 additions and 104 deletions.
diff --git a/doc/conf.py b/doc/conf.py
@@ -32,8 +32,16 @@
     "sphinx.ext.napoleon",
     "sphinx.ext.linkcode",
     "sphinx.ext.mathjax",
+    "sphinx_design",
+    "sphinx.ext.intersphinx"
 ]
 
+intersphinx_mapping = {
+    "jax": ("https://jax.readthedocs.io/en/latest", None),
+    "numpy": ("https://numpy.org/doc/stable", None),
+    "torch": ("https://pytorch.org/docs/stable", None),
+}
+
 needs_sphinx = "3"
 
 todo_include_todos = True

diff --git a/doc/environment.yml b/doc/environment.yml
@@ -13,6 +13,7 @@ dependencies:
   - mock
   - pillow
   - pymc-sphinx-theme
+  - sphinx-design
   - pip
   - pip:
     - -e ..