ci: run smoke tests across a python x torch version matrix

.github/workflows/ci.yaml

@@ -38,11 +38,23 @@ jobs:
         run: make check
 
   test-smoke:
-    name: Linux / make test-smoke
+    name: Linux / make test-smoke / torch=${{ matrix.torch_version }}
     runs-on: ubuntu-latest
     timeout-minutes: 60
     env:
       INSTALL_PRECOMMIT: "false"
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - torch_group: torch_2_8
+            torch_version: "2.8"
+          - torch_group: torch_2_9
+            torch_version: "2.9"
+          - torch_group: torch_2_10
+            torch_version: "2.10"
+          - torch_group: torch_2_11
+            torch_version: "2.11"
     steps:
       - name: Check out repository
         uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
@@ -51,7 +63,7 @@ jobs:
         with:
           enable-cache: false
       - name: Run `make test-smoke`
-        run: make test-smoke
+        run: make test-smoke TORCH_GROUP=${{ matrix.torch_group }}
 
   test-tutorials:
     name: Linux / make test-tutorials

Makefile

@@ -53,6 +53,11 @@ VENV_HIGHEST_TORCH ?= .venv-highest-torch
 VENV_LOWEST_TORCH ?= .venv-lowest-torch
 VENV_TUTORIAL ?= .venv-tutorial
 
+# Torch dependency group (pyproject.toml [dependency-groups]) that every
+# environment-building target (env, test, test-smoke, docs, ...) pins to.
+TORCH_GROUP ?= torch_2_11
+export TORCH_GROUP
+
 # Documentation directory. Defaults to $(MAKEFILE_DIR)docs so the same recipe
 # works in both contexts:
 #
@@ -176,7 +181,7 @@ env: _maybe_patch_pyproject
 
 # Set up development environment with latest supported PyTorch version
 env-highest-torch: _maybe_patch_pyproject
-	@$(SETUP_ENV) --venv $(VENV_HIGHEST_TORCH) --python-version $(PYTHON_VERSION) --with-highest_tested_torch
+	@TORCH_GROUP=torch_2_11 $(SETUP_ENV) --venv $(VENV_HIGHEST_TORCH) --python-version $(PYTHON_VERSION)
 	@$(call write_active_venv,$(VENV_HIGHEST_TORCH))
 
 # Set up environment for running tutorials (quantization notebook)
@@ -235,6 +240,7 @@ test-slow:
 	@$(MAKE) test PYTEST_ARGS="--marker slow"
 
 # Run smoke tests only (pass PYTEST_ARGS for custom flags, e.g., make test-smoke PYTEST_ARGS="--junitxml=results.xml").
+# Pass TORCH_GROUP to smoke test against a specific torch version (default: torch_2_11).
 test-smoke:
 	@$(call use_env,VENV) && \
 	echo "Running smoke tests..." && \
@@ -244,20 +250,20 @@ test-smoke:
 # Run tests on lowest supported PyTorch version (pass PYTEST_ARGS for custom flags)
 test-lowest-pytorch:
 	@echo "Running tests on lowest PyTorch version supported..."
-	@$(call use_env,VENV_LOWEST_TORCH,--with-lowest_tested_torch) && \
-		echo "Testing with lowest supported PyTorch versions" && \
-		uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
-		$(RUN_TESTS) $(PYTEST_ARGS) && \
-		echo "All tests passed!"
+	@TORCH_GROUP=torch_2_8 $(call use_env,VENV_LOWEST_TORCH) && \
+	echo "Testing with lowest supported PyTorch versions" && \
+	uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
+	$(RUN_TESTS) $(PYTEST_ARGS) && \
+	echo "All tests passed!"
 
 # Run tests on highest supported PyTorch version (pass PYTEST_ARGS for custom flags)
 test-highest-pytorch:
 	@echo "Running tests on highest PyTorch version supported..."
-	@$(call use_env,VENV_HIGHEST_TORCH,--with-highest_tested_torch) && \
-		echo "Testing with latest supported PyTorch versions" && \
-		uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
-		$(RUN_TESTS) $(PYTEST_ARGS) && \
-		echo "All tests passed!"
+	@TORCH_GROUP=torch_2_11 $(call use_env,VENV_HIGHEST_TORCH) && \
+	echo "Testing with latest supported PyTorch versions" && \
+	uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
+	$(RUN_TESTS) $(PYTEST_ARGS) && \
+	echo "All tests passed!"
 
 # Run tutorial notebook tests
 test-tutorials:

ci/nox/noxfile.py

@@ -34,8 +34,14 @@
 options.default_venv_backend = "uv"
 options.error_on_missing_interpreters = True
 
+TORCH_GROUP = os.environ.get("TORCH_GROUP")
 
-@session(python=get_supported_python_versions(), uv_extras=["coreai"], uv_groups=["test"])
+
+@session(
+    python=get_supported_python_versions(),
+    uv_extras=["coreai"],
+    uv_groups=["test", TORCH_GROUP],
+)
 def smoke_tests(session: Session) -> None:
     """Smoke test the package build and coreai_opt imports and basic functionality.
 

pyproject.toml

@@ -46,7 +46,7 @@ dependencies = [
     # based on other package versions. We can either 1) add a stricter check to require the newer torchao
     # version for everyone, or 2) add a runtime check in src/coreai_opt/init.py to ensure torchao
     # version >= 0.15.0 for torch version >= 2.9.0. Opting option 1.
-    # These torch versions must be in bounds of lowest_tested_torch and highest_tested_torch
+    # These torch versions must be in bounds of torch_2_8, torch_2_9, torch_2_10, and torch_2_11
     "torch>=2.8.0,<=2.11.0",
     "torchao>=0.15.0,<=0.17.0",
     "tqdm>=4.65",
@@ -107,20 +107,6 @@ docs = [
 # This applies to the `coreml` group as well.
 coreai = [ "coreai-opt[coreai]" ]
 coreml = [ "coreai-opt[coreml]" ]
-# Used in CI to force latest mimimum supported torch version
-# These torch versions must be in bounds of torch versions listed in project dependencies
-highest_tested_torch = [
-    "torch==2.11.0",
-    "torchao==0.17.0",
-    "torchvision==0.26.0",
-]
-# Used in CI to force lowest mimimum supported torch version
-# These torch versions must be in bounds of torch versions listed in project dependencies
-lowest_tested_torch = [
-    "torch==2.8.0",
-    "torchao==0.15.0",
-    "torchvision==0.23.0",
-]
 pre-commit = [
     "bashate>=2.1.1",
     "darker>=3.0.0",
@@ -144,6 +130,29 @@ pre-commit = [
     "tomli-w>=1.0.0",
 ]
 tamm-export = []
+# Used in CI/smoke tests to pin specific torch minor versions for the
+# cartesian python x torch smoke test matrix.
+# These torch versions must be in bounds of torch versions listed in project dependencies
+torch_2_10 = [
+    "torch==2.10.0",
+    "torchao==0.16.0",
+    "torchvision==0.25.0",
+]
+torch_2_11 = [
+    "torch==2.11.0",
+    "torchao==0.17.0",
+    "torchvision==0.26.0",
+]
+torch_2_8 = [
+    "torch==2.8.0",
+    "torchao==0.15.0",
+    "torchvision==0.23.0",
+]
+torch_2_9 = [
+    "torch==2.9.1",
+    "torchao==0.15.0",
+    "torchvision==0.24.1",
+]
 # Since torch and torchao are project dependencies, we need to include torchvision in dev
 # This allows standard `make test` to find torchvision
 torchvision = [
@@ -190,8 +199,10 @@ exclude-newer-package = { coreai-core = false, coreai-torch = false }
 # Declare conflicting groups so uv does not error
 conflicts = [
     [
-        { group = "highest_tested_torch" },
-        { group = "lowest_tested_torch" },
+        { group = "torch_2_8" },
+        { group = "torch_2_9" },
+        { group = "torch_2_10" },
+        { group = "torch_2_11" },
     ],
 ]
 [tool.uv.sources]

scripts/make/setup_env.sh

@@ -59,8 +59,8 @@ fi
 #    - in_section && /^\[/: If we're in the section AND hit another section header (line starting with [)
 #      * exit: Stop processing (we've left the dependency-groups section)
 #    - in_section {print}: If we're in the section, print the line
-# 2. grep -E '^[a-z_-]+ = \[': Filter to lines that define groups
-#    - ^[a-z_-]+: Group name at start of line (lowercase letters, hyphens, underscores)
+# 2. grep -E '^[a-z0-9_-]+ = \[': Filter to lines that define groups
+#    - ^[a-z0-9_-]+: Group name at start of line (lowercase letters, digits, hyphens, underscores)
 #    - = \[: Followed by space, equals sign, space, opening bracket
 # 3. cut -d' ' -f1: Extract just the group name
 #    - -d' ': Use space as delimiter
@@ -75,7 +75,7 @@ AVAILABLE_GROUPS=$(
         in_section && /^\[/ { exit }
         in_section { print }
     ' "$PYPROJECT_TOML" |
-        grep -E '^[a-z_-]+ = \[' |
+        grep -E '^[a-z0-9_-]+ = \[' |
         cut -d' ' -f1 |
         tr '\n' ' '
 )
@@ -109,7 +109,7 @@ ENSURE_MODE=false
 
 # Groups excluded from --all-groups due to mutual conflicts in pyproject.toml.
 # tamm-export is omitted because it's opt-in only (never in default-groups or --all-groups).
-CONFLICTING_GROUPS=("highest_tested_torch" "lowest_tested_torch")
+CONFLICTING_GROUPS=("torch_2_8" "torch_2_9" "torch_2_10" "torch_2_11")
 
 show_help() {
     echo "Usage: $0 [OPTIONS]"
@@ -127,6 +127,7 @@ show_help() {
     echo ""
     echo "Environment variables:"
     echo "  VENV                Virtual environment name, overrides --venv (default: .venv)"
+    echo "  TORCH_GROUP         Torch dependency group to pin; source of truth for torch pinning"
     echo ""
     echo "Available dependency groups: $AVAILABLE_GROUPS"
     echo "Conflicting groups (excluded from --all-groups): ${CONFLICTING_GROUPS[*]}"
@@ -135,8 +136,8 @@ show_help() {
     echo "  $0 --python-version 3.11                               # Setup with dev group only"
     echo "  $0 --python-version 3.11 --with-docs                   # Setup with dev and docs groups"
     echo "  $0 --python-version 3.11 --all-groups                                    # Setup with all non-conflicting groups"
-    echo "  $0 --python-version 3.11 --all-groups --with-highest_tested_torch        # Setup with all groups and highest torch"
-    echo "  $0 --python-version 3.11 --all-groups --with-lowest_tested_torch         # Setup with all groups and lowest torch"
+    echo "  TORCH_GROUP=torch_2_11 $0 --python-version 3.11 --all-groups             # Setup with all groups and torch 2.11"
+    echo "  TORCH_GROUP=torch_2_8 $0 --python-version 3.11 --all-groups              # Setup with all groups and torch 2.8"
     echo "  $0 --python-version 3.11 --venv .venv-exp              # Setup with custom venv name"
     echo "  $0 --python-version 3.11 --with-docs --venv .venv-exp  # Setup with docs group and custom venv name"
     echo "  $0 --python-version 3.12                               # Setup with Python 3.12"
@@ -222,6 +223,7 @@ validate_groups() {
 # Validate dependency group names early
 [[ ${#EXTRA_GROUPS[@]} -gt 0 ]] && validate_groups "${EXTRA_GROUPS[@]}"
 [[ ${#EXCLUDE_GROUPS[@]} -gt 0 ]] && validate_groups "${EXCLUDE_GROUPS[@]}"
+[[ -n "${TORCH_GROUP:-}" ]] && validate_groups "$TORCH_GROUP"
 
 # Check for conflicts between --with-<group> and --without-<group>
 if [[ ${#EXTRA_GROUPS[@]} -gt 0 && ${#EXCLUDE_GROUPS[@]} -gt 0 ]]; then
@@ -267,24 +269,35 @@ if [[ "$ENSURE_MODE" == "true" ]] && [ -f "$VENV/bin/python" ]; then
         for GROUP in "${EXTRA_GROUPS[@]}"; do
             case "$GROUP" in
             docs) IMPORT_STMTS+="; import sphinx" ;;
-            highest_tested_torch | lowest_tested_torch)
-                IMPORT_STMTS+="; import torchao"
-                EXPECTED_TORCH="$(group_torch_pin "$GROUP")"
-                # These groups always pin torch, so an empty result means the
-                # pyproject parse regressed — fail loudly instead of silently
-                # skipping the version check (which would reintroduce the bug).
-                if [[ -z "$EXPECTED_TORCH" ]]; then
-                    echo "Error: could not parse a torch pin for group '$GROUP' in $PYPROJECT_TOML" >&2
-                    exit 1
-                fi
-                IMPORT_STMTS+="; import torch; assert torch.__version__.split('+')[0] == '$EXPECTED_TORCH'"
-                ;;
             rio) IMPORT_STMTS+="; import turi_lightning" ;;
             tamm-export) IMPORT_STMTS+="; import tamm_export" ;;
             esac
         done
     fi
 
+    # TORCH_GROUP is the single source of truth for torch pinning; fall back
+    # to a --with-torch_2_* flag for direct, non-Make invocations. Checked
+    # unconditionally rather than only when present in EXTRA_GROUPS, since
+    # TORCH_GROUP drives the sync regardless of what's in EXTRA_GROUPS.
+    EXPECTED_GROUP="${TORCH_GROUP:-}"
+    if [[ -z "$EXPECTED_GROUP" ]]; then
+        for GROUP in "${EXTRA_GROUPS[@]:-}"; do
+            [[ " ${CONFLICTING_GROUPS[*]} " == *" ${GROUP} "* ]] && EXPECTED_GROUP="$GROUP"
+        done
+    fi
+    if [[ -n "$EXPECTED_GROUP" ]]; then
+        IMPORT_STMTS+="; import torchao"
+        EXPECTED_TORCH="$(group_torch_pin "$EXPECTED_GROUP")"
+        # These groups always pin torch, so an empty result means the
+        # pyproject parse regressed — fail loudly instead of silently
+        # skipping the version check (which would reintroduce the bug).
+        if [[ -z "$EXPECTED_TORCH" ]]; then
+            echo "Error: could not parse a torch pin for '$EXPECTED_GROUP' in $PYPROJECT_TOML" >&2
+            exit 1
+        fi
+        IMPORT_STMTS+="; import torch; assert torch.__version__.split('+')[0] == '$EXPECTED_TORCH'"
+    fi
+
     if "$VENV/bin/python" -c "$IMPORT_STMTS" 2>/dev/null; then
         exit 0
     fi
@@ -324,9 +337,13 @@ echo "[2/3] Installing dependencies..."
 SYNC_CMD=(uv sync --active)
 if [[ "$ALL_GROUPS" == "true" ]]; then
     SYNC_CMD+=(--all-groups)
-    # Exclude conflicting groups unless explicitly requested via --with-*
+    # Exclude conflicting torch groups other than the one TORCH_GROUP or
+    # --with-<group> selects, so --all-groups doesn't try to sync every
+    # torch_2_* group at once. If both TORCH_GROUP and --with-<other-group>
+    # name different groups, neither gets excluded here and `uv sync` itself
+    # rejects the combination via its own conflicting-groups resolution.
     for GROUP in "${CONFLICTING_GROUPS[@]}"; do
-        if [[ ! " ${EXTRA_GROUPS[*]:-} " == *" ${GROUP} "* ]]; then
+        if [[ "$GROUP" != "${TORCH_GROUP:-}" ]] && [[ ! " ${EXTRA_GROUPS[*]:-} " == *" ${GROUP} "* ]]; then
             SYNC_CMD+=(--no-group "$GROUP")
         fi
     done
@@ -335,6 +352,14 @@ elif [[ ${#EXTRA_GROUPS[@]} -gt 0 ]]; then
         SYNC_CMD+=(--group "$GROUP")
     done
 fi
+# TORCH_GROUP is the single source of truth for torch pinning: append it
+# unconditionally (already covered under --all-groups by not being
+# excluded above). A disagreeing --with-<other-torch-group> ends up as a
+# second, different --group flag here, which `uv sync` itself rejects via
+# its conflicting-groups resolution rather than a bespoke check.
+if [[ -n "${TORCH_GROUP:-}" && "$ALL_GROUPS" != "true" ]]; then
+    SYNC_CMD+=(--group "$TORCH_GROUP")
+fi
 # Apply explicit group exclusions (e.g., --without-coreai)
 if [[ ${#EXCLUDE_GROUPS[@]} -gt 0 ]]; then
     for GROUP in "${EXCLUDE_GROUPS[@]}"; do