5959# - in_section && /^\[/: If we're in the section AND hit another section header (line starting with [)
6060# * exit: Stop processing (we've left the dependency-groups section)
6161# - in_section {print}: If we're in the section, print the line
62- # 2. grep -E '^[a-z_ -]+ = \[': Filter to lines that define groups
63- # - ^[a-z_- ]+: Group name at start of line (lowercase letters, hyphens, underscores)
62+ # 2. grep -E '^[a-z0-9_ -]+ = \[': Filter to lines that define groups
63+ # - ^[a-z0-9_- ]+: Group name at start of line (lowercase letters, digits , hyphens, underscores)
6464# - = \[: Followed by space, equals sign, space, opening bracket
6565# 3. cut -d' ' -f1: Extract just the group name
6666# - -d' ': Use space as delimiter
@@ -75,7 +75,7 @@ AVAILABLE_GROUPS=$(
7575 in_section && /^\[/ { exit }
7676 in_section { print }
7777 ' " $PYPROJECT_TOML " |
78- grep -E ' ^[a-z_ -]+ = \[' |
78+ grep -E ' ^[a-z0-9_ -]+ = \[' |
7979 cut -d' ' |
8080 tr ' \n' ' '
8181)
@@ -109,7 +109,7 @@ ENSURE_MODE=false
109109
110110# Groups excluded from --all-groups due to mutual conflicts in pyproject.toml.
111111# tamm-export is omitted because it's opt-in only (never in default-groups or --all-groups).
112- CONFLICTING_GROUPS=(" highest_tested_torch " " lowest_tested_torch "
112+ CONFLICTING_GROUPS=(" torch_2_8 " " torch_2_9 " " torch_2_10 " " torch_2_11 "
113113
114114show_help () {
115115 echo " Usage: $0 [OPTIONS]"
@@ -127,6 +127,7 @@ show_help() {
127127 echo " "
128128 echo " Environment variables:"
129129 echo " VENV Virtual environment name, overrides --venv (default: .venv)"
130+ echo " TORCH_GROUP Torch dependency group to pin; source of truth for torch pinning"
130131 echo " "
131132 echo " Available dependency groups: $AVAILABLE_GROUPS "
132133 echo " Conflicting groups (excluded from --all-groups): ${CONFLICTING_GROUPS[*]} "
@@ -135,8 +136,8 @@ show_help() {
135136 echo " $0 --python-version 3.11 # Setup with dev group only"
136137 echo " $0 --python-version 3.11 --with-docs # Setup with dev and docs groups"
137138 echo " $0 --python-version 3.11 --all-groups # Setup with all non-conflicting groups"
138- echo " $0 --python-version 3.11 --all-groups --with-highest_tested_torch # Setup with all groups and highest torch"
139- echo " $0 --python-version 3.11 --all-groups --with-lowest_tested_torch # Setup with all groups and lowest torch"
139+ echo " TORCH_GROUP=torch_2_11 $0 --python-version 3.11 --all-groups # Setup with all groups and torch 2.11 "
140+ echo " TORCH_GROUP=torch_2_8 $0 --python-version 3.11 --all-groups # Setup with all groups and torch 2.8 "
140141 echo " $0 --python-version 3.11 --venv .venv-exp # Setup with custom venv name"
141142 echo " $0 --python-version 3.11 --with-docs --venv .venv-exp # Setup with docs group and custom venv name"
142143 echo " $0 --python-version 3.12 # Setup with Python 3.12"
@@ -222,6 +223,7 @@ validate_groups() {
222223# Validate dependency group names early
223224[[ ${# EXTRA_GROUPS[@]} -gt 0 ]] && validate_groups " ${EXTRA_GROUPS[@]} "
224225[[ ${# EXCLUDE_GROUPS[@]} -gt 0 ]] && validate_groups " ${EXCLUDE_GROUPS[@]} "
226+ [[ -n " ${TORCH_GROUP:- } " && validate_groups " $TORCH_GROUP "
225227
226228# Check for conflicts between --with-<group> and --without-<group>
227229if [[ ${# EXTRA_GROUPS[@]} -gt 0 && ${# EXCLUDE_GROUPS[@]} -gt 0 ]]; then
@@ -267,24 +269,35 @@ if [[ "$ENSURE_MODE" == "true" ]] && [ -f "$VENV/bin/python" ]; then
267269 for GROUP in " ${EXTRA_GROUPS[@]} " ; do
268270 case " $GROUP " in
269271 docs) IMPORT_STMTS+=" ; import sphinx"
270- highest_tested_torch | lowest_tested_torch)
271- IMPORT_STMTS+=" ; import torchao"
272- EXPECTED_TORCH=" $( group_torch_pin " $GROUP " ) "
273- # These groups always pin torch, so an empty result means the
274- # pyproject parse regressed — fail loudly instead of silently
275- # skipping the version check (which would reintroduce the bug).
276- if [[ -z " $EXPECTED_TORCH " ; then
277- echo " Error: could not parse a torch pin for group '$GROUP ' in $PYPROJECT_TOML " >&2
278- exit 1
279- fi
280- IMPORT_STMTS+=" ; import torch; assert torch.__version__.split('+')[0] == '$EXPECTED_TORCH '"
281- ;;
282272 rio) IMPORT_STMTS+=" ; import turi_lightning"
283273 tamm-export) IMPORT_STMTS+=" ; import tamm_export"
284274 esac
285275 done
286276 fi
287277
278+ # TORCH_GROUP is the single source of truth for torch pinning; fall back
279+ # to a --with-torch_2_* flag for direct, non-Make invocations. Checked
280+ # unconditionally rather than only when present in EXTRA_GROUPS, since
281+ # TORCH_GROUP drives the sync regardless of what's in EXTRA_GROUPS.
282+ EXPECTED_GROUP=" ${TORCH_GROUP:- } "
283+ if [[ -z " $EXPECTED_GROUP " ; then
284+ for GROUP in " ${EXTRA_GROUPS[@]:- } " ; do
285+ [[ " ${CONFLICTING_GROUPS[*]} " == * " ${GROUP} " * ]] && EXPECTED_GROUP=" $GROUP "
286+ done
287+ fi
288+ if [[ -n " $EXPECTED_GROUP " ; then
289+ IMPORT_STMTS+=" ; import torchao"
290+ EXPECTED_TORCH=" $( group_torch_pin " $EXPECTED_GROUP " ) "
291+ # These groups always pin torch, so an empty result means the
292+ # pyproject parse regressed — fail loudly instead of silently
293+ # skipping the version check (which would reintroduce the bug).
294+ if [[ -z " $EXPECTED_TORCH " ; then
295+ echo " Error: could not parse a torch pin for '$EXPECTED_GROUP ' in $PYPROJECT_TOML " >&2
296+ exit 1
297+ fi
298+ IMPORT_STMTS+=" ; import torch; assert torch.__version__.split('+')[0] == '$EXPECTED_TORCH '"
299+ fi
300+
288301 if " $VENV /bin/python" " $IMPORT_STMTS " 2> /dev/null; then
289302 exit 0
290303 fi
@@ -324,9 +337,13 @@ echo "[2/3] Installing dependencies..."
324337SYNC_CMD=(uv sync --active)
325338if [[ " $ALL_GROUPS " == " true" ; then
326339 SYNC_CMD+=(--all-groups)
327- # Exclude conflicting groups unless explicitly requested via --with-*
340+ # Exclude conflicting torch groups other than the one TORCH_GROUP or
341+ # --with-<group> selects, so --all-groups doesn't try to sync every
342+ # torch_2_* group at once. If both TORCH_GROUP and --with-<other-group>
343+ # name different groups, neither gets excluded here and `uv sync` itself
344+ # rejects the combination via its own conflicting-groups resolution.
328345 for GROUP in " ${CONFLICTING_GROUPS[@]} " ; do
329- if [[ ! " ${EXTRA_GROUPS[*]:- } " == * " ${GROUP} " * ]]; then
346+ if [[ " $GROUP " != " ${TORCH_GROUP :- } " ]] && [[ ! " ${EXTRA_GROUPS[*]:- } " == * " ${GROUP} " * ]]; then
330347 SYNC_CMD+=(--no-group " $GROUP "
331348 fi
332349 done
@@ -335,6 +352,14 @@ elif [[ ${#EXTRA_GROUPS[@]} -gt 0 ]]; then
335352 SYNC_CMD+=(--group " $GROUP "
336353 done
337354fi
355+ # TORCH_GROUP is the single source of truth for torch pinning: append it
356+ # unconditionally (already covered under --all-groups by not being
357+ # excluded above). A disagreeing --with-<other-torch-group> ends up as a
358+ # second, different --group flag here, which `uv sync` itself rejects via
359+ # its conflicting-groups resolution rather than a bespoke check.
360+ if [[ -n " ${TORCH_GROUP:- } " && " $ALL_GROUPS " != " true" ; then
361+ SYNC_CMD+=(--group " $TORCH_GROUP "
362+ fi
338363# Apply explicit group exclusions (e.g., --without-coreai)
339364if [[ ${# EXCLUDE_GROUPS[@]} -gt 0 ]]; then
340365 for GROUP in " ${EXCLUDE_GROUPS[@]} " ; do
0 commit comments