Check all compatible torch indexes when --torch-backend is enabled

crates/uv-client/src/registry_client.rs

@@ -258,11 +258,27 @@ impl RegistryClient {
     fn index_urls_for(&self, package_name: &PackageName) -> impl Iterator<Item = &IndexUrl> {
         self.torch_backend
             .as_ref()
-            .and_then(|torch_backend| torch_backend.index_urls(package_name))
+            .and_then(|torch_backend| {
+                torch_backend
+                    .applies_to(package_name)
+                    .then(|| torch_backend.index_urls())
+            })
             .map(Either::Left)
             .unwrap_or_else(|| Either::Right(self.index_urls.indexes().map(Index::url)))
     }
 
+    /// Return the appropriate [`IndexStrategy`] for the given [`PackageName`].
+    fn index_strategy_for(&self, package_name: &PackageName) -> IndexStrategy {
+        self.torch_backend
+            .as_ref()
+            .and_then(|torch_backend| {
+                torch_backend
+                    .applies_to(package_name)
+                    .then_some(IndexStrategy::UnsafeFirstMatch)
+            })
+            .unwrap_or(self.index_strategy)
+    }
+
     /// Fetch a package from the `PyPI` simple API.
     ///
     /// "simple" here refers to [PEP 503 – Simple Repository API](https://peps.python.org/pep-0503/)
@@ -290,7 +306,7 @@ impl RegistryClient {
 
         let mut results = Vec::new();
 
-        match self.index_strategy {
+        match self.index_strategy_for(package_name) {
             // If we're searching for the first index that contains the package, fetch serially.
             IndexStrategy::FirstIndex => {
                 for index in indexes {

crates/uv-torch/src/backend.rs

@@ -157,12 +157,9 @@ impl TorchStrategy {
         }
     }
 
-    /// Return the appropriate index URLs for the given [`TorchStrategy`] and [`PackageName`].
-    pub fn index_urls(
-        &self,
-        package_name: &PackageName,
-    ) -> Option<impl Iterator<Item = &IndexUrl>> {
-        if !matches!(
+    /// Returns `true` if the [`TorchStrategy`] applies to the given [`PackageName`].
+    pub fn applies_to(&self, package_name: &PackageName) -> bool {
+        matches!(
             package_name.as_str(),
             "torch"
                 | "torch-model-archiver"
@@ -176,32 +173,31 @@ impl TorchStrategy {
                 | "torchtext"
                 | "torchvision"
                 | "pytorch-triton"
-        ) {
-            return None;
-        }
+        )
+    }
 
+    /// Return the appropriate index URLs for the given [`TorchStrategy`].
+    pub fn index_urls(&self) -> impl Iterator<Item = &IndexUrl> {
         match self {
             TorchStrategy::Auto { os, driver_version } => {
                 // If this is a GPU-enabled package, and CUDA drivers are installed, use PyTorch's CUDA
                 // indexes.
                 //
                 // See: https://github.com/pmeier/light-the-torch/blob/33397cbe45d07b51ad8ee76b004571a4c236e37f/light_the_torch/_patch.py#L36-L49
                 match os {
-                    Os::Manylinux { .. } | Os::Musllinux { .. } => {
-                        Some(Either::Left(Either::Left(
-                            LINUX_DRIVERS
-                                .iter()
-                                .filter_map(move |(backend, version)| {
-                                    if driver_version >= version {
-                                        Some(backend.index_url())
-                                    } else {
-                                        None
-                                    }
-                                })
-                                .chain(std::iter::once(TorchBackend::Cpu.index_url())),
-                        )))
-                    }
-                    Os::Windows => Some(Either::Left(Either::Right(
+                    Os::Manylinux { .. } | Os::Musllinux { .. } => Either::Left(Either::Left(
+                        LINUX_DRIVERS
+                            .iter()
+                            .filter_map(move |(backend, version)| {
+                                if driver_version >= version {
+                                    Some(backend.index_url())
+                                } else {
+                                    None
+                                }
+                            })
+                            .chain(std::iter::once(TorchBackend::Cpu.index_url())),
+                    )),
+                    Os::Windows => Either::Left(Either::Right(
                         WINDOWS_CUDA_VERSIONS
                             .iter()
                             .filter_map(move |(backend, version)| {
@@ -212,22 +208,20 @@ impl TorchStrategy {
                                 }
                             })
                             .chain(std::iter::once(TorchBackend::Cpu.index_url())),
-                    ))),
+                    )),
                     Os::Macos { .. }
                     | Os::FreeBsd { .. }
                     | Os::NetBsd { .. }
                     | Os::OpenBsd { .. }
                     | Os::Dragonfly { .. }
                     | Os::Illumos { .. }
                     | Os::Haiku { .. }
-                    | Os::Android { .. } => Some(Either::Right(std::iter::once(
-                        TorchBackend::Cpu.index_url(),
-                    ))),
+                    | Os::Android { .. } => {
+                        Either::Right(std::iter::once(TorchBackend::Cpu.index_url()))
+                    }
                 }
             }
-            TorchStrategy::Backend(backend) => {
-                Some(Either::Right(std::iter::once(backend.index_url())))
-            }
+            TorchStrategy::Backend(backend) => Either::Right(std::iter::once(backend.index_url())),
         }
     }
 }