feat: make TORCH_GROUP the single source of truth for torch pinning

setup_env.sh now reads TORCH_GROUP directly and folds it into every
sync it runs, so make env/test/test-fast/docs/etc. all pin torch
instead of floating on whatever pyproject.toml's plain bound resolves
to. test-highest-pytorch/test-lowest-pytorch/env-highest-torch lock
their group via an inline `export TORCH_GROUP=... &&` at the top of
the recipe (rather than a hardcoded --with-torch_2_X flag), so their
name stays a promise regardless of a command-line override — verified
that GNU Make's own `override` directive achieves the same lock but
trips up the mbake Makefile formatter's duplicate-target detector,
which doesn't recognize that syntax.

--with-torch_2_X stays as a guarded alternate path for direct,
non-Make invocations of setup_env.sh; a guard now errors clearly if it
disagrees with an explicitly-set TORCH_GROUP instead of surfacing a
raw uv conflicting-groups resolver error.

Author: guru-desh

Makefile

@@ -53,10 +53,15 @@ VENV_HIGHEST_TORCH ?= .venv-highest-torch
 VENV_LOWEST_TORCH ?= .venv-lowest-torch
 VENV_TUTORIAL ?= .venv-tutorial
 
-# Torch dependency group (pyproject.toml [dependency-groups]) that
-# `test-smoke` installs; override to smoke test against a different pinned
-# torch minor version, e.g. `make test-smoke TORCH_GROUP=torch_2_8`.
+# Torch dependency group (pyproject.toml [dependency-groups]) that every
+# environment-building target (env, test, test-smoke, docs, ...) pins to.
+# Single source of truth for which torch/torchao/torchvision triple
+# setup_env.sh installs; override e.g. `make test-fast TORCH_GROUP=torch_2_8`.
+# test-highest-pytorch/test-lowest-pytorch/env-highest-torch lock this to a
+# specific group via `export TORCH_GROUP=... &&` at the top of their recipe,
+# so their name stays a promise regardless of a command-line value.
 TORCH_GROUP ?= torch_2_11
+export TORCH_GROUP
 
 # Documentation directory. Defaults to $(MAKEFILE_DIR)docs so the same recipe
 # works in both contexts:
@@ -181,7 +186,8 @@ 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-torch_2_11
+	@export 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)
@@ -250,20 +256,22 @@ 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-torch_2_8) && \
-		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!"
+	@export 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-torch_2_11) && \
-		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!"
+	@export 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:

scripts/make/setup_env.sh

@@ -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-torch_2_11                   # Setup with all groups and torch 2.11"
-    echo "  $0 --python-version 3.11 --all-groups --with-torch_2_8                    # Setup with all groups and torch 2.8"
+    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
@@ -233,6 +235,21 @@ if [[ ${#EXTRA_GROUPS[@]} -gt 0 && ${#EXCLUDE_GROUPS[@]} -gt 0 ]]; then
     done
 fi
 
+# TORCH_GROUP (env var) is the single source of truth for torch pinning; a
+# --with-<torch_2_*> flag that disagrees with it is almost certainly a
+# mistake (e.g. TORCH_GROUP left set from a prior `make` invocation) rather
+# than an intentional double-specification, so fail loudly instead of
+# leaving it to a raw `uv` conflicting-groups resolver error.
+if [[ -n "${TORCH_GROUP:-}" ]]; then
+    for GROUP in "${EXTRA_GROUPS[@]:-}"; do
+        if [[ " ${CONFLICTING_GROUPS[*]} " == *" ${GROUP} "* && "$GROUP" != "$TORCH_GROUP" ]]; then
+            echo "Error: --with-$GROUP conflicts with TORCH_GROUP=$TORCH_GROUP." >&2
+            echo "TORCH_GROUP is the source of truth; unset it or set it to '$GROUP' instead." >&2
+            exit 1
+        fi
+    done
+fi
+
 # Validate venv name contains only allowed characters
 if [[ ! "$VENV" =~ ^[a-zA-Z0-9._/-]+$ ]]; then
     echo "Error: Virtual environment name contains invalid characters"
@@ -267,24 +284,35 @@ if [[ "$ENSURE_MODE" == "true" ]] && [ -f "$VENV/bin/python" ]; then
         for GROUP in "${EXTRA_GROUPS[@]}"; do
             case "$GROUP" in
             docs) IMPORT_STMTS+="; import sphinx" ;;
-            torch_2_8 | torch_2_9 | torch_2_10 | torch_2_11)
-                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 +352,14 @@ 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 trip the mutual
+    # -conflicts declaration in pyproject.toml.
     for GROUP in "${CONFLICTING_GROUPS[@]}"; do
-        if [[ ! " ${EXTRA_GROUPS[*]:-} " == *" ${GROUP} "* ]]; then
+        GROUP_REQUESTED=false
+        [[ "$GROUP" == "${TORCH_GROUP:-}" ]] && GROUP_REQUESTED=true
+        [[ " ${EXTRA_GROUPS[*]:-} " == *" ${GROUP} "* ]] && GROUP_REQUESTED=true
+        if [[ "$GROUP_REQUESTED" == "false" ]]; then
             SYNC_CMD+=(--no-group "$GROUP")
         fi
     done
@@ -335,6 +368,15 @@ elif [[ ${#EXTRA_GROUPS[@]} -gt 0 ]]; then
         SYNC_CMD+=(--group "$GROUP")
     done
 fi
+# TORCH_GROUP is the single source of truth for torch pinning: fold it in
+# unless --all-groups already swept it up or a --with-<group> above already
+# requested it explicitly (the earlier guard ensures no disagreement).
+TORCH_GROUP_ALREADY_SYNCED=false
+[[ "$ALL_GROUPS" == "true" ]] && TORCH_GROUP_ALREADY_SYNCED=true
+[[ " ${EXTRA_GROUPS[*]:-} " == *" ${TORCH_GROUP:-} "* ]] && TORCH_GROUP_ALREADY_SYNCED=true
+if [[ -n "${TORCH_GROUP:-}" && "$TORCH_GROUP_ALREADY_SYNCED" == "false" ]]; 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