From 22aa2a5951eeaa5e58aae684c5d054263ab15e4c Mon Sep 17 00:00:00 2001 From: laurent Date: Thu, 27 Jul 2023 09:26:00 +0100 Subject: [PATCH] Do not panic on empty ranges. --- candle-core/src/indexer.rs | 2 +- candle-core/tests/indexing_tests.rs | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/candle-core/src/indexer.rs b/candle-core/src/indexer.rs index c7a9a140b1..6ab209ca31 100644 --- a/candle-core/src/indexer.rs +++ b/candle-core/src/indexer.rs @@ -42,7 +42,7 @@ impl Tensor { Bound::Excluded(n) => *n, Bound::Unbounded => dims[i], }; - let out = x.narrow(current_dim, start, stop - start)?; + let out = x.narrow(current_dim, start, stop.saturating_sub(start))?; current_dim += 1; out } diff --git a/candle-core/tests/indexing_tests.rs b/candle-core/tests/indexing_tests.rs index 6a71b8fb98..122f63b05a 100644 --- a/candle-core/tests/indexing_tests.rs +++ b/candle-core/tests/indexing_tests.rs @@ -58,6 +58,19 @@ fn range_index() -> Result<()> { let result = tensor.i(..=1)?; assert_eq!(result.dims(), &[2, 3]); assert_eq!(result.to_vec2::()?, &[[0, 1, 2], [3, 4, 5]]); + + // Empty range + let result = tensor.i(1..1)?; + assert_eq!(result.dims(), &[0, 3]); + let empty: [[u32; 3]; 0] = []; + assert_eq!(result.to_vec2::()?, &empty); + + // Similar to PyTorch, allow empty ranges when the computed length is negative. + #[allow(clippy::reversed_empty_ranges)] + let result = tensor.i(1..0)?; + assert_eq!(result.dims(), &[0, 3]); + let empty: [[u32; 3]; 0] = []; + assert_eq!(result.to_vec2::()?, &empty); Ok(()) }