Fix up

Author: charliermarsh

Cargo.lock

@@ -63,6 +63,7 @@ uv-small-str = { path = "crates/uv-small-str" }
 uv-state = { path = "crates/uv-state" }
 uv-static = { path = "crates/uv-static" }
 uv-tool = { path = "crates/uv-tool" }
+uv-torch = { path = "crates/uv-torch" }
 uv-trampoline-builder = { path = "crates/uv-trampoline-builder" }
 uv-types = { path = "crates/uv-types" }
 uv-version = { path = "crates/uv-version" }

Cargo.toml

@@ -28,6 +28,7 @@ uv-python = { workspace = true, features = ["clap", "schemars"]}
 uv-resolver = { workspace = true, features = ["clap"] }
 uv-settings = { workspace = true, features = ["schemars"] }
 uv-static = { workspace = true }
+uv-torch = { workspace = true, features = ["clap"] }
 uv-version = { workspace = true }
 uv-warnings = { workspace = true }
 

crates/uv-cli/Cargo.toml

@@ -21,6 +21,7 @@ use uv_pypi_types::VerbatimParsedUrl;
 use uv_python::{PythonDownloads, PythonPreference, PythonVersion};
 use uv_resolver::{AnnotationStyle, ExcludeNewer, ForkStrategy, PrereleaseMode, ResolutionMode};
 use uv_static::EnvVars;
+use uv_torch::TorchMode;
 
 pub mod comma;
 pub mod compat;
@@ -1258,6 +1259,19 @@ pub struct PipCompileArgs {
     #[arg(long, overrides_with("emit_index_annotation"), hide = true)]
     pub no_emit_index_annotation: bool,
 
+    /// The backend to use when fetching packages in the `PyTorch` ecosystem (e.g., `cu126` or `auto`)
+    ///
+    /// When set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem,
+    /// and will instead use the defined backend.
+    ///
+    /// For example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`,
+    /// uv will use the `PyTorch` index for CUDA 12.6.
+    ///
+    /// The `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently
+    /// installed CUDA drivers.
+    #[arg(long, value_enum)]
+    pub torch_backend: Option<TorchMode>,
+
     #[command(flatten)]
     pub compat_args: compat::PipCompileCompatArgs,
 }
@@ -1499,6 +1513,19 @@ pub struct PipSyncArgs {
     #[arg(long)]
     pub dry_run: bool,
 
+    /// The backend to use when fetching packages in the `PyTorch` ecosystem (e.g., `cu126` or `auto`)
+    ///
+    /// When set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem,
+    /// and will instead use the defined backend.
+    ///
+    /// For example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`,
+    /// uv will use the `PyTorch` index for CUDA 12.6.
+    ///
+    /// The `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently
+    /// installed CUDA drivers.
+    #[arg(long, value_enum)]
+    pub torch_backend: Option<TorchMode>,
+
     #[command(flatten)]
     pub compat_args: compat::PipSyncCompatArgs,
 }
@@ -1791,6 +1818,19 @@ pub struct PipInstallArgs {
     #[arg(long)]
     pub dry_run: bool,
 
+    /// The backend to use when fetching packages in the `PyTorch` ecosystem (e.g., `cu126` or `auto`)
+    ///
+    /// When set, uv will ignore the configured index URLs for packages in the `PyTorch` ecosystem,
+    /// and will instead use the defined backend.
+    ///
+    /// For example, when set to `cpu`, uv will use the CPU-only `PyTorch` index; when set to `cu126`,
+    /// uv will use the `PyTorch` index for CUDA 12.6.
+    ///
+    /// The `auto` mode will attempt to detect the appropriate `PyTorch` index based on the currently
+    /// installed CUDA drivers.
+    #[arg(long, value_enum)]
+    pub torch_backend: Option<TorchMode>,
+
     #[command(flatten)]
     pub compat_args: compat::PipInstallCompatArgs,
 }

crates/uv-cli/src/lib.rs

@@ -25,6 +25,7 @@ uv-platform-tags = { workspace = true }
 uv-pypi-types = { workspace = true }
 uv-small-str = { workspace = true }
 uv-static = { workspace = true }
+uv-torch = { workspace = true }
 uv-version = { workspace = true }
 uv-warnings = { workspace = true }