[XLA:Python] Fix DLPack behavior with unit dimensions.

As discovered in jax-ml/jax#24680, when a PyTorch tensor has a dimension with size `1`, it seems to report the DLPack stride for that dimension as `1`. This behavior wasn't supported by the logic in XLA, resulting in an incorrect layout on the imported array.

PiperOrigin-RevId: 696341186

xla/python/dlpack.cc

@@ -218,14 +218,17 @@ absl::StatusOr<std::vector<int64_t>> StridesToLayout(
   std::vector<int64_t> minor_to_major(dims.size());
   std::iota(minor_to_major.begin(), minor_to_major.end(), 0);
   absl::c_sort(minor_to_major, [&](int a, int b) {
+    // If two dimensions have the same stride, prefer the major-to-minor
+    // interpretation of the ordering, since that's what JAX wants.
+    if (dims[a] <= 1 || dims[b] <= 1) {
+      return b < a;
+    }
     if (strides[a] < strides[b]) {
       return true;
     }
     if (strides[a] > strides[b]) {
       return false;
     }
-    // If two dimensions have the same stride, prefer the major-to-minor
-    // interpretation of the ordering, since that's what JAX wants.
     return b < a;
   });
   int64_t stride = 1;