From de633c452c5d9df870dfb1df7e736b08efc18b58 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Tue, 25 Aug 2026 14:59:25 +0200 Subject: [PATCH] Fix slice to support empty range and ranges that end at 0. --- .../tensor/tensorops/StructuralOps.scala | 8 +++- .../tensor/TensorOpsStructureSuite.scala | 38 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala index fd41342c..12ef5760 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala @@ -534,7 +534,13 @@ object StructuralOps: case AxisAtIndex(_, idx) => indicesBuffer(dimIndex) = py.Any.from(idx) case AxisAtRange(_, range) => - indicesBuffer(dimIndex) = PySlice(range.head, range.last + 1, range.step) + require(range.isEmpty || (range.min >= 0 && range.max < dimSize), s"$range is out of bounds for axis of size $dimSize") + // map Scala Range to Python Range which is exclusive + val stop = range match + case r: Range.Inclusive => r.end + r.step.sign + case r: Range.Exclusive => r.end + // the range is in bounds, so a negative stop can only mean "before the first element", which Python spells as None + indicesBuffer(dimIndex) = PySlice(range.start, if stop < 0 then py.None else py.Any.from(stop), range.step) case AxisAtIndices(_, indices) => indicesBuffer(dimIndex) = indices.map(py.Any.from).toPythonCopy // TODO find out why Copy is needed here case AxisAtTupleIndices(_, indices) => diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index 6b84bb28..7674eccc 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala @@ -457,6 +457,44 @@ class TensorOpsStructureSuite extends DimwitTest: val t3 = t.set(Axis[A].at(1 to 2))(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f))) t3 should approxEqual(Tensor1(Axis[A]).fromArray(Array(1.0f, 7.0f, 8.0f))) + it("slice at a range"): + val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f)) + t.slice(Axis[A].at(1 to 2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(2.0f, 3.0f))) + t.slice(Axis[A].at(1 until 3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(2.0f, 3.0f))) + t.slice(Axis[A].at(0 until 4 by 2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(1.0f, 3.0f))) + + it("slice at a descending range"): + val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f)) + t.slice(Axis[A].at(3 to 1 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f))) + t.slice(Axis[A].at(3 until 0 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f))) + // both spellings of a descending range down to index 0 + t.slice(Axis[A].at(3 to 0 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f, 1.0f))) + t.slice(Axis[A].at(3 until -1 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f, 1.0f))) + t.slice(Axis[A].at(3 to 0 by -2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 2.0f))) + + it("slice rejects a range that leaves the axis"): + val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f)) + the[IllegalArgumentException] thrownBy t.slice(Axis[A].at(3 until -10000 by -1)) should have message + "requirement failed: Range 3 until -10000 by -1 is out of bounds for axis of size 4" + an[IllegalArgumentException] should be thrownBy t.slice(Axis[A].at(3 until -3 by -1)) + an[IllegalArgumentException] should be thrownBy t.slice(Axis[A].at(0 until 10)) + // a very negative end is fine as long as the step steps over the invalid indices + t.slice(Axis[A].at(3 until -3 by -3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 1.0f))) + + it("slice at an empty range yields a zero-length axis"): + val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) + val empty = t.slice(Axis[A].at(0 until 0)) + empty.axes shouldBe List("A", "B") + empty.shape(Axis[A]) shouldBe 0 + empty.shape(Axis[B]) shouldBe 2 + t.slice(Axis[A].at(2 until 2)).shape(Axis[A]) shouldBe 0 + t.slice(Axis[A].at(1 until 1)).shape(Axis[A]) shouldBe 0 + + it("set at an empty range is a no-op"): + val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) + val empty = Tensor1(Axis[A]).fromArray(Array.empty[Float]) + t.set(Axis[A].at(1 until 1))(empty) should approxEqual(t) + describe("AxisAtIndices"): it("set at a seq"):