Remove ModelWithAux in trainstep
#2435
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Unittests & Auto-publish | |
| # Allow to trigger the workflow manually (e.g. when deps changes) | |
| on: [push, workflow_dispatch] | |
| jobs: | |
| pytest-job: | |
| runs-on: ubuntu-latest | |
| timeout-minutes: 30 | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.ref }} | |
| cancel-in-progress: true | |
| steps: | |
| - uses: actions/checkout@v3 | |
| # Install deps | |
| - uses: actions/setup-python@v4 | |
| with: | |
| python-version: "3.11" | |
| # Uncomment to cache of pip dependencies (if tests too slow) | |
| # cache: pip | |
| # cache-dependency-path: '**/pyproject.toml' | |
| - run: pip --version | |
| # TODO(epot): Delete once all projects have been released | |
| - run: pip install git+https://github.com/google/etils | |
| - run: pip install git+https://github.com/google/flax | |
| - run: pip install -e .[dev] | |
| - run: pip freeze | |
| # Run tests (in parallel) | |
| # Filter out: | |
| # * Projects: Not part of core Kauldron (could be tested separately) | |
| # * TF Data pipeline (not supported due to TFGrain not open-sourced) | |
| # * XManager tests (not yet supported) | |
| # * sweep_utils_test: Depends on kxm | |
| # * lpips_test: Missing VGG weights | |
| # * partial_loader_test: Orbax partial checkpoint loader not yet open-sourced (TODO(epot): Restore) | |
| # * typing tests: Not yet supported due to typeguard version issues. | |
| - name: Run core tests | |
| run: | | |
| pytest -vv -n auto \ | |
| --ignore=projects/ \ | |
| --ignore=kauldron/data/tf/ \ | |
| --ignore=kauldron/xm/ \ | |
| --ignore=kauldron/metrics/lpips_test.py \ | |
| --ignore=kauldron/checkpoints/partial_loader_test.py \ | |
| --ignore=kauldron/utils/sweep_utils_test.py \ | |
| --ignore=kauldron/typing/shape_spec_test.py \ | |
| --ignore=kauldron/typing/type_check_test.py | |
| # Auto-publish when version is increased | |
| publish-job: | |
| # Only try to publish if: | |
| # * Repo is self (prevents running from forks) | |
| # * Branch is `main` | |
| if: | | |
| github.repository == 'google-research/kauldron' | |
| && github.ref == 'refs/heads/main' | |
| needs: pytest-job # Only publish after tests are successful | |
| runs-on: ubuntu-latest | |
| permissions: | |
| contents: write | |
| timeout-minutes: 30 | |
| steps: | |
| # Publish the package (if local `__version__` > pip version) | |
| - uses: etils-actions/pypi-auto-publish@v1 | |
| with: | |
| pypi-token: ${{ secrets.PYPI_API_TOKEN }} | |
| gh-token: ${{ secrets.GITHUB_TOKEN }} | |
| parse-changelog: true |