refactor: align argument names with PyTorch

.github/workflows/lint.yml

@@ -42,13 +42,13 @@ jobs:
       - run: |
           CUDA_VERSION="${{steps.cuda-toolkit.outputs.cuda}}"
           echo "CUDA_VERSION=${CUDA_VERSION}" >> "${GITHUB_ENV}"
-          TORCH_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
-          echo "TORCH_INDEX_URL=${TORCH_INDEX_URL}" >> "${GITHUB_ENV}"
+          PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
+          echo "PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}" >> "${GITHUB_ENV}"
 
           echo "Installed CUDA version is: ${CUDA_VERSION}"
           echo "CUDA install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
           nvcc -V
-          echo "Torch index URL: ${TORCH_INDEX_URL}"
+          echo "Torch index URL: ${PIP_EXTRA_INDEX_URL}"
 
       - name: Upgrade pip
         run: |
@@ -92,8 +92,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          python -m pip install --extra-index-url "${TORCH_INDEX_URL}" \
-            -r docs/requirements.txt
+          python -m pip install -r docs/requirements.txt
 
       - name: docstyle
         run: |

.github/workflows/tests.yml

@@ -53,25 +53,31 @@ jobs:
       - run: |
           CUDA_VERSION="${{steps.cuda-toolkit.outputs.cuda}}"
           echo "CUDA_VERSION=${CUDA_VERSION}" >> "${GITHUB_ENV}"
-          TORCH_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
-          echo "TORCH_INDEX_URL=${TORCH_INDEX_URL}" >> "${GITHUB_ENV}"
+          PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
+          echo "PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}" >> "${GITHUB_ENV}"
 
           echo "Installed CUDA version is: ${CUDA_VERSION}"
           echo "CUDA install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
           nvcc -V
-          echo "Torch index URL: ${TORCH_INDEX_URL}"
+          echo "Torch index URL: ${PIP_EXTRA_INDEX_URL}"
 
       - name: Upgrade pip
         run: |
           python -m pip install --upgrade pip setuptools
 
+      - name: Install PyTorch and FuncTorch nightly
+        run: |
+          export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL//whl/whl\/nightly}"
+          python -m pip install 'torch >= 1.13.0dev' ninja
+          python -m pip install git+https://github.com/pytorch/functorch.git
+
       - name: Install dependencies
         run: |
-          python -m pip install --extra-index-url "${TORCH_INDEX_URL}" \
-            -r tests/requirements.txt
+          python -m pip install -r tests/requirements.txt
 
       - name: Install TorchOpt
         run: |
+          export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL//whl/whl\/nightly}"
           python -m pip install -vvv -e .
 
       - name: Test with pytest

.gitignore

@@ -394,12 +394,12 @@ fabric.properties
 
 ##### Vim.gitignore #####
 # Swap
-[._]*.s[a-v][a-z]
+.*.s[a-v][a-z]
 !*.svg  # comment out if you don't need vector files
-[._]*.sw[a-p]
-[._]s[a-rt-v][a-z]
-[._]ss[a-gi-z]
-[._]sw[a-p]
+.*.sw[a-p]
+.s[a-rt-v][a-z]
+.ss[a-gi-z]
+.sw[a-p]
 
 # Session
 Session.vim

CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add optional argument `params` to update function in gradient transformations by [@XuehaiPan](https://github.com/XuehaiPan) in [#65](https://github.com/metaopt/torchopt/pull/65).
+- Add option `weight_decay` option to optimizers by [@XuehaiPan](https://github.com/XuehaiPan) in [#65](https://github.com/metaopt/torchopt/pull/65).
 - Add option `maximize` option to optimizers by [@XuehaiPan](https://github.com/XuehaiPan) in [#64](https://github.com/metaopt/torchopt/pull/64).
 - Refactor tests using `pytest.mark.parametrize` and enabling parallel testing by [@XuehaiPan](https://github.com/XuehaiPan) and [@Benjamin-eecs](https://github.com/Benjamin-eecs) in [#55](https://github.com/metaopt/torchopt/pull/55).
 - Add maml-omniglot few-shot classification example using functorch.vmap by [@Benjamin-eecs](https://github.com/Benjamin-eecs) in [#39](https://github.com/metaopt/torchopt/pull/39).
@@ -21,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Align argument names with PyTorch by [@XuehaiPan](https://github.com/XuehaiPan) in [#65](https://github.com/metaopt/torchopt/pull/65).
 - Replace JAX PyTrees with OpTree by [@XuehaiPan](https://github.com/XuehaiPan) in [#62](https://github.com/metaopt/torchopt/pull/62).
 - Update image link in README to support PyPI rendering by [@Benjamin-eecs](https://github.com/Benjamin-eecs) in [#56](https://github.com/metaopt/torchopt/pull/56).
 

Dockerfile

@@ -31,15 +31,15 @@ RUN echo "export PS1='[\[\e[1;33m\]\u\[\e[0m\]:\[\e[1;35m\]\w\[\e[0m\]]\$ '" >>
 
 # Setup virtual environment
 RUN /usr/bin/python3.9 -m venv --upgrade-deps ~/venv && rm -rf ~/.pip/cache
-RUN TORCH_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')" && \
-    echo "export TORCH_INDEX_URL='${TORCH_INDEX_URL}'" >> ~/venv/bin/activate && \
+RUN PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')" && \
+    echo "export PIP_EXTRA_INDEX_URL='${PIP_EXTRA_INDEX_URL}'" >> ~/venv/bin/activate && \
     echo "source /home/torchopt/venv/bin/activate" >> ~/.bashrc
 
 # Install dependencies
 WORKDIR /home/torchopt/torchopt
 COPY --chown=torchopt requirements.txt requirements.txt
 RUN source ~/venv/bin/activate && \
-    python -m pip install --extra-index-url "${TORCH_INDEX_URL}" -r requirements.txt && \
+    python -m pip install -r requirements.txt && \
     rm -rf ~/.pip/cache ~/.cache/pip
 
 ####################################################################################################
@@ -63,8 +63,7 @@ RUN go install github.com/google/addlicense@latest
 COPY --chown=torchopt tests/requirements.txt tests/requirements.txt
 COPY --chown=torchopt tutorials/requirements.txt tutorials/requirements.txt
 RUN source ~/venv/bin/activate && \
-    python -m pip install --extra-index-url "${TORCH_INDEX_URL}" \
-        -r tests/requirements.txt -r tutorials/requirements.txt && \
+    python -m pip install -r tests/requirements.txt -r tutorials/requirements.txt && \
     rm -rf ~/.pip/cache ~/.cache/pip
 
 ####################################################################################################

tests/helpers.py

@@ -26,8 +26,8 @@
 from torch.utils import data
 
 
-BATCH_SIZE = 4
-NUM_UPDATES = 3
+BATCH_SIZE = 64
+NUM_UPDATES = 5
 
 MODEL_NUM_INPUTS = 28 * 28  # MNIST
 MODEL_NUM_CLASSES = 10
@@ -82,13 +82,23 @@ def get_models(
             bias=True,
             dtype=dtype,
         ),
+        nn.BatchNorm1d(
+            num_features=MODEL_HIDDEN_SIZE,
+            track_running_stats=True,
+            dtype=dtype,
+        ),
         nn.ReLU(),
         nn.Linear(
             in_features=MODEL_HIDDEN_SIZE,
             out_features=MODEL_HIDDEN_SIZE,
             bias=True,
             dtype=dtype,
         ),
+        nn.BatchNorm1d(
+            num_features=MODEL_HIDDEN_SIZE,
+            track_running_stats=True,
+            dtype=dtype,
+        ),
         nn.ReLU(),
         nn.Linear(
             in_features=MODEL_HIDDEN_SIZE,
@@ -99,7 +109,7 @@ def get_models(
         nn.Softmax(dim=-1),
     )
     for name, param in model_base.named_parameters(recurse=True):
-        if name.endswith('weight'):
+        if name.endswith('weight') and param.ndim >= 2:
             nn.init.orthogonal_(param)
         if name.endswith('bias'):
             param.data.normal_(0, 0.1)
@@ -160,6 +170,13 @@ def assert_all_close(
         actual = actual - base
         expected = expected - base
 
+    if rtol is None or atol is None:
+        from torch.testing._comparison import get_tolerances
+
+        rtol, atol = get_tolerances(actual, expected, rtol=rtol, atol=atol)
+        rtol *= 4 * NUM_UPDATES
+        atol *= 4 * NUM_UPDATES
+
     torch.testing.assert_close(
         actual,
         expected,