Merge branch 'main' into fix-implicit-oop

XuehaiPan · XuehaiPan · commit e1bef1f8d104 · 2022-11-21T14:03:03.000+08:00
diff --git a/Makefile b/Makefile
@@ -10,6 +10,7 @@ COMMIT_HASH    = $(shell git log -1 --format=%h)
 PATH           := $(HOME)/go/bin:$(PATH)
 PYTHON         ?= $(shell command -v python3 || command -v python)
 CLANG_FORMAT   ?= $(shell command -v clang-format-14 || command -v clang-format)
+PYTESTOPTS     ?=
 
 .PHONY: default
 default: install
@@ -104,7 +105,7 @@ pytest: pytest-install
 	cd tests && \
 	$(PYTHON) -m pytest --verbose --color=yes --durations=0 \
 		--cov="$(PROJECT_NAME)" --cov-report=xml --cov-report=term-missing \
-		.
+		$(PYTESTOPTS) .
 
 test: pytest
 
diff --git a/README.md b/README.md
@@ -212,7 +212,7 @@ def stationary(params, meta_params, data):
     return stationary condition
 
 # Decorator for wrapping the function
-# and specify the linear solver (conjugate gradient or Neumann series)
+# Optionally specify the linear solver (conjugate gradient or Neumann series)
 @torchopt.diff.implicit.custom_root(stationary, solve=linear_solver)
 def solve(params, meta_params, data):
     # Forward optimization process for params
@@ -233,29 +233,29 @@ Users need to define the stationary condition/objective function and the inner-l
 
 ```python
 # Inherited from the class ImplicitMetaGradientModule
-# and specify the linear solver (conjugate gradient or Neumann series)
+# Optionally specify the linear solver (conjugate gradient or Neumann series)
 class InnerNet(ImplicitMetaGradientModule, linear_solver):
     def __init__(self, meta_param):
         super().__init__()
         self.meta_param = meta_param
         ...
 
-    def forward(self, data):
+    def forward(self, batch):
         # Forward process
         ...
 
-    def optimality(self, data):
+    def optimality(self, batch, labels):
         # Stationary condition construction for calculating implicit gradient
         # NOTE: If this method is not implemented, it will be automatically
         # derived from the gradient of the `objective` function.
         ...
 
-    def objective(self, data):
+    def objective(self, batch, labels):
         # Define the inner-loop optimization objective
         ...
 
-    def solve(self, data):
-        # conduct the inner-loop optimization
+    def solve(self, batch, labels):
+        # Conduct the inner-loop optimization
         ...
 
 # Get meta_params and data
@@ -282,12 +282,12 @@ Refer to the tutorial notebook [Zero-order Differentiation](tutorials/6_Zero_Ord
 
 ```python
 # Customize the noise sampling function in ES
-def sample(params, batch, labels, *, sample_shape):
+def sample(sample_shape):
     ...
     return sample_noise
 
 # Specify method and hyper-parameter of ES
-@torchopt.diff.zero_order(method, sample)
+@torchopt.diff.zero_order(sample, method)
 def forward(params, batch, labels):
     # forward process
     return output
diff --git a/conda-recipe.yaml b/conda-recipe.yaml
@@ -42,7 +42,6 @@ dependencies:
   - optax                    # for tutorials
   - jaxopt                   # for tests
   - tensorboard              # for examples
-  - wandb
 
   # Device select
   - nvidia/label/cuda-11.7.1::cuda-toolkit = 11.7
@@ -68,13 +67,13 @@ dependencies:
   - setproctitle
 
   # Documentation
-  - sphinx
+  - sphinx >= 5.2.1
   - sphinx_rtd_theme
   - sphinx-autobuild
   - sphinx-copybutton
   - sphinxcontrib-spelling
   - sphinxcontrib-bibtex
-  - sphinx-autodoc-typehints
+  - sphinx-autodoc-typehints >= 1.19.2
   - pyenchant
   - hunspell-en
   - myst-nb
diff --git a/docs/conda-recipe.yaml b/docs/conda-recipe.yaml
@@ -57,13 +57,13 @@ dependencies:
   - pillow
 
   # Documentation
-  - sphinx
+  - sphinx >= 5.2.1
   - sphinx_rtd_theme
   - sphinx-autobuild
   - sphinx-copybutton
   - sphinxcontrib-spelling
   - sphinxcontrib-bibtex
-  - sphinx-autodoc-typehints
+  - sphinx-autodoc-typehints >= 1.19.2
   - pyenchant
   - hunspell-en
   - myst-nb
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -154,7 +154,7 @@ def filter(self, record):
 
 # See: https://sphinxcontrib-katex.readthedocs.io/en/0.4.1/macros.html
 latex_macros = r"""
-    \def \d              #1{\operatorname{#1}}
+    \def \d #1{\operatorname{#1}}
 """
 
 # Translate LaTeX macros to KaTeX and add to options for HTML builder
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -23,10 +23,10 @@ endif()
 
 list(APPEND torchopt_csrc "${adam_op_src}")
 
-pybind11_add_module(_C THIN_LTO "${torchopt_csrc}")
+pybind11_add_module(_C MODULE THIN_LTO "${torchopt_csrc}")
 
 target_link_libraries(
     _C PRIVATE
-    ${TORCH_LIBRARIES}
+    "${TORCH_LIBRARIES}"
     OpenMP::OpenMP_CXX
 )
