Forces u to be strictly positive.

huggingface · Aug 17, 2024 · 8a6b355 · 8a6b355
1 parent 2936f84
commit 8a6b355
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/candle-core/src/shape.rs b/candle-core/src/shape.rs
@@ -329,7 +329,7 @@ impl Dim for D {
         match self {
             Self::Minus1 if rank >= 1 => Ok(rank - 1),
             Self::Minus2 if rank >= 2 => Ok(rank - 2),
-            Self::Minus(u) if rank >= *u => Ok(rank - *u),
+            Self::Minus(u) if *u > 0 && rank >= *u => Ok(rank - *u),
             _ => Err(self.out_of_range(shape, op)),
         }
     }
@@ -339,7 +339,7 @@ impl Dim for D {
         match self {
             Self::Minus1 => Ok(rank),
             Self::Minus2 if rank >= 1 => Ok(rank - 1),
-            Self::Minus(u) if rank + 1 >= *u => Ok(rank + 1 - *u),
+            Self::Minus(u) if *u > 0 && rank + 1 >= *u => Ok(rank + 1 - *u),
             _ => Err(self.out_of_range(shape, op)),
         }
     }