From 191ea86bb8022611a264e7dfb4c484fd4bc71f5a Mon Sep 17 00:00:00 2001 From: shivasankar Date: Fri, 8 May 2026 00:06:11 +0900 Subject: [PATCH 1/4] add scalar overload for slice setitem --- numojo/core/ndarray.mojo | 165 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 165 insertions(+) diff --git a/numojo/core/ndarray.mojo b/numojo/core/ndarray.mojo index 8556c747..1b5120b8 100644 --- a/numojo/core/ndarray.mojo +++ b/numojo/core/ndarray.mojo @@ -2504,6 +2504,171 @@ struct NDArray[dtype: DType = DType.float64]( self.__setitem__(slices=slice_list, val=val) + def _setitem_slice_scalar( + mut self, slices: List[Slice], val: Scalar[Self.dtype] + ) raises: + """Internal backend: fill every element in a slice region with a scalar. + + Called by the user-facing `__setitem__(*Slice, scalar=)` and + `__setitem__(*Variant[Slice,Int], scalar=)` overloads after they + normalise their arguments into a `List[Slice]`. + + Args: + slices: One `Slice` per array dimension (trailing dims already + padded to full range by the caller). + val: The scalar value to write to every selected position. + """ + var n_slices: Int = len(slices) + var slice_list: List[InternalSlice] = self._adjust_slice(slices) + + for i in range(n_slices, self.ndim): + slice_list.append(InternalSlice(0, self.shape[i], 1)) + + # Use a single flat counter + coordinate reconstruction to stay + # allocation-free and layout-agnostic (works for C, F, strided). + var total: Int = 1 + var region_shape = List[Int](capacity=self.ndim) + for i in range(self.ndim): + var slen: Int + var s = slice_list[i] + if s.step > 0: + slen = max((s.end - s.start + s.step - 1) // s.step, 0) + else: + slen = max((s.start - s.end - s.step - 1) // (-s.step), 0) + region_shape.append(slen) + total *= slen + + var coords = List[Int](capacity=self.ndim) + for _ in range(self.ndim): + coords.append(0) + + for _ in range(total): + var buf_idx: Int = self.offset + for d in range(self.ndim): + buf_idx += ( + slice_list[d].start + coords[d] * slice_list[d].step + ) * self.strides[d] + self._buf.ptr.store(buf_idx, val) + + for d in range(self.ndim - 1, -1, -1): + coords[d] += 1 + if coords[d] < region_shape[d]: + break + coords[d] = 0 + + def __setitem__( + mut self, *slices: Slice, scalar: Scalar[Self.dtype] + ) raises: + """Sets all elements in the slice region to a scalar value. + + Delegates to `_setitem_slice_scalar` after packing slices into a list. + + Note: the trailing keyword is named `scalar` (not `val`) to avoid + ambiguity with the `*slices: Slice, val: Self` overload during Mojo + overload resolution. + + Args: + slices: Variadic slices, one per dimension (trailing dims default + to the full range). + scalar: The scalar value to broadcast into every selected position. + + Examples: + ```mojo + var a = nm.arange[nm.i32](16).reshape(nm.Shape(4, 4)) + a.__setitem__(Slice(1,3), Slice(1,3), scalar=99) + ``` + """ + var slice_list = List[Slice](capacity=slices.__len__()) + for i in range(slices.__len__()): + slice_list.append(slices[i]) + self._setitem_slice_scalar(slice_list, scalar) + + def __setitem__( + mut self, *slices: Variant[Slice, Int], val: Scalar[Self.dtype] + ) raises: + """Sets elements selected by mixed integer/slice indices to a scalar. + + Handles two cases: + - All entries are integers (full coordinate) → direct single-element + write via `_setitem`, no allocation. + - Mixed or slice-only → normalise integers to unit-length slices and + delegate to `__setitem__(List[Slice], Scalar)` for the region fill. + + Note: the trailing keyword is named `scalar` (not `val`) to avoid + ambiguity with the `*slices: Variant[Slice, Int], val: Self` overload + during Mojo overload resolution. + + Args: + slices: Variadic mix of `Slice` and `Int` index entries. + scalar: The scalar value to write. + + Examples: + ```mojo + var a = nm.arange[nm.i32](16).reshape(nm.Shape(4, 4)) + a.__setitem__(1, 2, scalar=99) # single element + a.__setitem__(1, Slice(2,4), scalar=0) # one row, partial column + a.__setitem__(Slice(1,3), Slice(2,4), scalar=7) # sub-matrix + ``` + """ + var n = slices.__len__() + if n > self.ndim: + raise Error( + NumojoError( + category="index", + message=String( + "Too many indices: received {} but array has only {}" + " dimensions." + ).format(n, self.ndim), + location=( + "NDArray.__setitem__(*slices: Variant[Slice, Int]," + " scalar: Scalar)" + ), + ) + ) + + var slice_list = List[Slice](capacity=self.ndim) + var count_int: Int = 0 + var coords = List[Int]() + + # Track which array dimension each entry maps to (Int and Slice both + # consume one dimension; the loop index equals the array dim here + # because Variant[Slice,Int] carries no NewAxis). + for i in range(n): + if slices[i].isa[Int](): + var idx = slices[i][Int] + if idx >= self.shape[i] or idx < -self.shape[i]: + raise Error( + NumojoError( + category="index", + message=String( + "Integer index {} out of bounds for axis {}" + " (size {}). Valid range: [-{}, {})." + ).format(idx, i, self.shape[i], self.shape[i], self.shape[i]), + location=( + "NDArray.__setitem__(*slices: Variant[Slice," + " Int], scalar: Scalar)" + ), + ) + ) + if idx < 0: + idx += self.shape[i] + count_int += 1 + coords.append(idx) + slice_list.append(Slice(idx, idx + 1, 1)) + else: + slice_list.append(slices[i][Slice]) + + # Fast path: every dimension was given an integer → single element. + if count_int == self.ndim: + self.itemset(coords.copy(), scalar) + return + + # Pad trailing dimensions. + for i in range(n, self.ndim): + slice_list.append(Slice(0, self.shape[i], 1)) + + self._setitem_slice_scalar(slice_list, scalar) + def __setitem__( mut self, index: NDArray[DType.int], val: NDArray[Self.dtype] ) raises: From 648d1ce11ca6c6afe7c4007d3190afd4cc971423 Mon Sep 17 00:00:00 2001 From: shivasankar Date: Fri, 8 May 2026 00:06:15 +0900 Subject: [PATCH 2/4] add tests --- tests/core/test_setitem_scalar.mojo | 309 ++++++++++++++++++++++++++++ 1 file changed, 309 insertions(+) create mode 100644 tests/core/test_setitem_scalar.mojo diff --git a/tests/core/test_setitem_scalar.mojo b/tests/core/test_setitem_scalar.mojo new file mode 100644 index 00000000..b81b94ad --- /dev/null +++ b/tests/core/test_setitem_scalar.mojo @@ -0,0 +1,309 @@ +from std.python import Python, PythonObject +from std.testing import TestSuite +from std.testing.testing import assert_equal, assert_true + +import numojo as nm +from numojo.prelude import * +from utils_for_test import check + + +# Mojo's __setitem__ overload resolution with variadic+keyword args requires: +# - List[Slice] scalar: pass the list positionally, `val=` keyword works. +# - *Slice scalar: use explicit `.__setitem__(s1, s2, ..., scalar=v)`. +# - *Variant[Slice,Int] scalar: use explicit `.__setitem__(i, s, scalar=v)`. +# The `scalar` keyword name was chosen (instead of `val`) to avoid shadowing +# the existing `val: Self` overloads at the Mojo compiler level. + + +# ===== Step 3: __setitem__(List[Slice], val: Scalar) ===== + + +def test_setitem_list_slice_scalar_1d() raises: + """List[Slice] scalar backend: fills selected 1D range.""" + var a = nm.arange[nm.i32](0, 6, step=1) + var sl = List[Slice]() + sl.append(Slice(1, 4)) + a._setitem_slice_scalar(sl, Scalar[nm.i32](99)) + + assert_equal(Int(a.item(0)), 0) + assert_equal(Int(a.item(1)), 99) + assert_equal(Int(a.item(2)), 99) + assert_equal(Int(a.item(3)), 99) + assert_equal(Int(a.item(4)), 4) + assert_equal(Int(a.item(5)), 5) + + +def test_setitem_list_slice_scalar_2d_submatrix() raises: + """List[Slice] scalar: fills a 2x2 sub-matrix, surrounding values unchanged.""" + var a = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) + var sl = List[Slice]() + sl.append(Slice(1, 3)) + sl.append(Slice(1, 3)) + a._setitem_slice_scalar(sl, Scalar[nm.i32](0)) + + assert_equal(Int(a.item(1, 1)), 0) + assert_equal(Int(a.item(1, 2)), 0) + assert_equal(Int(a.item(2, 1)), 0) + assert_equal(Int(a.item(2, 2)), 0) + # Row 0 and col 0 untouched + assert_equal(Int(a.item(0, 1)), 1) + assert_equal(Int(a.item(1, 0)), 4) + assert_equal(Int(a.item(3, 3)), 15) + + +def test_setitem_list_slice_scalar_3d() raises: + """List[Slice] scalar: fills a 3D sub-region.""" + var a = nm.arange[nm.i32](0, 24, step=1).reshape(Shape(2, 3, 4)) + var sl = List[Slice]() + sl.append(Slice(0, 2)) + sl.append(Slice(1, 3)) + sl.append(Slice(2, 4)) + a._setitem_slice_scalar(sl, Scalar[nm.i32](5)) + + for i in range(2): + for j in range(1, 3): + for k in range(2, 4): + assert_equal( + Int(a.item(i, j, k)), + 5, + String("a[{},{},{}] should be 5").format(i, j, k), + ) + # Outside region + assert_equal(Int(a.item(0, 0, 0)), 0) + assert_equal(Int(a.item(0, 0, 1)), 1) + + +def test_setitem_list_slice_scalar_step() raises: + """List[Slice] scalar: step-stride fills every other row.""" + var a = nm.arange[nm.i32](0, 12, step=1).reshape(Shape(3, 4)) + var sl = List[Slice]() + sl.append(Slice(0, 3, 2)) # rows 0, 2 + sl.append(Slice(0, 4)) + a._setitem_slice_scalar(sl, Scalar[nm.i32](-1)) + + for c in range(4): + assert_equal(Int(a.item(0, c)), -1) + assert_equal(Int(a.item(2, c)), -1) + # Row 1 unchanged: 4 5 6 7 + assert_equal(Int(a.item(1, 0)), 4) + assert_equal(Int(a.item(1, 3)), 7) + + +def test_setitem_list_slice_scalar_implicit_trailing_dim() raises: + """List[Slice] scalar: fewer slices than ndim pads trailing dims as full.""" + var a = nm.arange[nm.i32](0, 12, step=1).reshape(Shape(3, 4)) + var sl = List[Slice]() + sl.append(Slice(1, 3)) # only row dim given + a._setitem_slice_scalar(sl, Scalar[nm.i32](42)) + + for c in range(4): + assert_equal(Int(a.item(1, c)), 42) + assert_equal(Int(a.item(2, c)), 42) + # Row 0 unchanged + for c in range(4): + assert_equal(Int(a.item(0, c)), c) + + +def test_setitem_list_slice_scalar_single_element() raises: + """List[Slice] scalar: unit-length slices write exactly one element.""" + var a = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) + var sl = List[Slice]() + sl.append(Slice(2, 3)) + sl.append(Slice(2, 3)) + a._setitem_slice_scalar(sl, Scalar[nm.i32](123)) + + assert_equal(Int(a.item(2, 2)), 123) + assert_equal(Int(a.item(2, 1)), 9) + assert_equal(Int(a.item(2, 3)), 11) + assert_equal(Int(a.item(1, 2)), 6) + assert_equal(Int(a.item(3, 2)), 14) + + +def test_setitem_list_slice_scalar_whole_array() raises: + """List[Slice] scalar: full-range slices zero the whole array.""" + var a = nm.arange[nm.i32](1, 10, step=1).reshape(Shape(3, 3)) + var sl = List[Slice]() + sl.append(Slice(0, 3)) + sl.append(Slice(0, 3)) + a._setitem_slice_scalar(sl, Scalar[nm.i32](0)) + + for i in range(3): + for j in range(3): + assert_equal(Int(a.item(i, j)), 0) + + +# ===== Step 3: __setitem__(*slices: Slice, scalar: Scalar) ===== + + +def test_setitem_variadic_slice_scalar_2d() raises: + """*Slice scalar wrapper produces same result as List[Slice] scalar.""" + var a = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) + var b = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) + + var sl = List[Slice]() + sl.append(Slice(1, 3)) + sl.append(Slice(1, 3)) + a._setitem_slice_scalar(sl, Scalar[nm.i32](77)) + + b.__setitem__(Slice(1, 3), Slice(1, 3), scalar=Scalar[nm.i32](77)) + + for i in range(4): + for j in range(4): + assert_equal( + Int(a.item(i, j)), + Int(b.item(i, j)), + String("*Slice and List[Slice] scalar paths diverge at ({},{})").format(i, j), + ) + + +def test_setitem_variadic_slice_scalar_1d() raises: + """*Slice scalar: 1D slice fill.""" + var a = nm.arange[nm.i32](0, 8, step=1) + a.__setitem__(Slice(2, 6), scalar=Scalar[nm.i32](5)) + + assert_equal(Int(a.item(0)), 0) + assert_equal(Int(a.item(1)), 1) + assert_equal(Int(a.item(2)), 5) + assert_equal(Int(a.item(5)), 5) + assert_equal(Int(a.item(6)), 6) + assert_equal(Int(a.item(7)), 7) + + +def test_setitem_variadic_slice_scalar_implicit_trailing() raises: + """*Slice scalar: single slice on 2D array, trailing dim implicit full.""" + var a = nm.arange[nm.i32](0, 12, step=1).reshape(Shape(3, 4)) + a.__setitem__(Slice(0, 2), scalar=Scalar[nm.i32](9)) + + for c in range(4): + assert_equal(Int(a.item(0, c)), 9) + assert_equal(Int(a.item(1, c)), 9) + # Row 2 unchanged: 8 9 10 11 + assert_equal(Int(a.item(2, 0)), 8) + assert_equal(Int(a.item(2, 3)), 11) + + +# ===== Step 4: __setitem__(*Variant[Slice, Int], scalar: Scalar) ===== + + +def test_setitem_mixed_single_element_2d() raises: + """All-Int fast path: single element write on 2D array.""" + var a = nm.arange[nm.i32](0, 12, step=1).reshape(Shape(3, 4)) + a.__setitem__(1, 2, scalar=Scalar[nm.i32](99)) + + assert_equal(Int(a.item(1, 2)), 99) + assert_equal(Int(a.item(1, 1)), 5) + assert_equal(Int(a.item(1, 3)), 7) + assert_equal(Int(a.item(0, 2)), 2) + assert_equal(Int(a.item(2, 2)), 10) + + +def test_setitem_mixed_single_element_3d() raises: + """All-Int fast path: single element write on 3D array.""" + var a = nm.arange[nm.i32](0, 24, step=1).reshape(Shape(2, 3, 4)) + a.__setitem__(0, 1, 2, scalar=Scalar[nm.i32](77)) + + assert_equal(Int(a.item(0, 1, 2)), 77) + assert_equal(Int(a.item(0, 0, 0)), 0) + assert_equal(Int(a.item(1, 2, 3)), 23) + + +def test_setitem_mixed_int_slice_row() raises: + """Int + Slice: select row by int, column range by slice.""" + var a = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) + a.__setitem__(1, Slice(1, 3), scalar=Scalar[nm.i32](55)) + + assert_equal(Int(a.item(1, 1)), 55) + assert_equal(Int(a.item(1, 2)), 55) + assert_equal(Int(a.item(1, 0)), 4) + assert_equal(Int(a.item(1, 3)), 7) + assert_equal(Int(a.item(0, 1)), 1) + assert_equal(Int(a.item(2, 1)), 9) + + +def test_setitem_mixed_slice_int_col() raises: + """Slice + Int: select row range by slice, column by int.""" + var a = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) + a.__setitem__(Slice(1, 3), 2, scalar=Scalar[nm.i32](33)) + + assert_equal(Int(a.item(1, 2)), 33) + assert_equal(Int(a.item(2, 2)), 33) + assert_equal(Int(a.item(1, 0)), 4) + assert_equal(Int(a.item(2, 3)), 11) + assert_equal(Int(a.item(0, 2)), 2) + assert_equal(Int(a.item(3, 2)), 14) + + +def test_setitem_mixed_negative_int() raises: + """Negative integer index is normalised correctly.""" + var a = nm.arange[nm.i32](0, 12, step=1).reshape(Shape(3, 4)) + a.__setitem__(-1, 2, scalar=Scalar[nm.i32](88)) + + assert_equal(Int(a.item(2, 2)), 88) + assert_equal(Int(a.item(0, 2)), 2) + assert_equal(Int(a.item(1, 2)), 6) + assert_equal(Int(a.item(2, 0)), 8) + + +def test_setitem_mixed_int_only_1d() raises: + """All-Int fast path on 1D array.""" + var a = nm.arange[nm.i32](0, 6, step=1) + a.__setitem__(3, scalar=Scalar[nm.i32](99)) + + assert_equal(Int(a.item(3)), 99) + assert_equal(Int(a.item(2)), 2) + assert_equal(Int(a.item(4)), 4) + + +def test_setitem_mixed_partial_indices_3d() raises: + """Two ints on 3D: trailing dim filled implicitly as full range.""" + var a = nm.arange[nm.i32](0, 24, step=1).reshape(Shape(2, 3, 4)) + a.__setitem__(0, 1, scalar=Scalar[nm.i32](11)) + + for k in range(4): + assert_equal( + Int(a.item(0, 1, k)), + 11, + String("a[0,1,{}] should be 11").format(k), + ) + assert_equal(Int(a.item(0, 0, 0)), 0) + assert_equal(Int(a.item(1, 1, 0)), 16) + + +def test_setitem_mixed_does_not_corrupt_neighbours() raises: + """Single-element write must not corrupt adjacent elements.""" + var a = nm.zeros[nm.i32](Shape(4, 4)) + a.__setitem__(1, 2, scalar=Scalar[nm.i32](7)) + + for r in range(4): + for c in range(4): + var expected = 7 if (r == 1 and c == 2) else 0 + assert_equal( + Int(a.item(r, c)), + expected, + String("Element ({},{}) wrong after a[1,2]=7").format(r, c), + ) + + +def test_setitem_mixed_agrees_with_list_slice_scalar() raises: + """Mixed int/slice scalar path matches List[Slice] scalar path.""" + var a = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) + var b = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) + + a.__setitem__(1, Slice(1, 3), scalar=Scalar[nm.i32](55)) + + var sl = List[Slice]() + sl.append(Slice(1, 2)) + sl.append(Slice(1, 3)) + b._setitem_slice_scalar(sl, Scalar[nm.i32](55)) + + for i in range(4): + for j in range(4): + assert_equal( + Int(a.item(i, j)), + Int(b.item(i, j)), + String("Mixed and List[Slice] scalar paths diverge at ({},{})").format(i, j), + ) + + +def main() raises: + TestSuite.discover_tests[__functions_in_module()]().run() From 7ac603239e245684533bab88e60c85d3a561fc60 Mon Sep 17 00:00:00 2001 From: shivasankar Date: Fri, 8 May 2026 00:53:55 +0900 Subject: [PATCH 3/4] fix keyword and clean up ndarry printing --- numojo/core/ndarray.mojo | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/numojo/core/ndarray.mojo b/numojo/core/ndarray.mojo index 1b5120b8..f7ceeec7 100644 --- a/numojo/core/ndarray.mojo +++ b/numojo/core/ndarray.mojo @@ -2584,7 +2584,7 @@ struct NDArray[dtype: DType = DType.float64]( self._setitem_slice_scalar(slice_list, scalar) def __setitem__( - mut self, *slices: Variant[Slice, Int], val: Scalar[Self.dtype] + mut self, *slices: Variant[Slice, Int], scalar: Scalar[Self.dtype] ) raises: """Sets elements selected by mixed integer/slice indices to a scalar. @@ -3865,20 +3865,25 @@ struct NDArray[dtype: DType = DType.float64]( ) else: try: + var order: String + if self.is_c_contiguous(): + order = "C" + elif self.is_f_contiguous(): + order = "F" + else: + order = "non-contiguous" writer.write( self._array_to_string(0, 0) + "\n" + String(self.ndim) - + "D-array Shape" - + String(self.shape) - + " Strides" - + String(self.strides) + + "D-array Shape: " + + self.shape.__str__() + + " Strides: " + + self.strides.__str__() + " DType: " + _concise_dtype_str(self.dtype) - + " C-cont: " - + String(self.is_c_contiguous()) - + " F-cont: " - + String(self.is_f_contiguous()) + + " order: " + + order + " own data: " + String(self.flags.OWNDATA) ) From 4bba516846716035111ac63931e98a4fc59bb948 Mon Sep 17 00:00:00 2001 From: shivasankar Date: Fri, 8 May 2026 01:09:34 +0900 Subject: [PATCH 4/4] fix format --- numojo/core/ndarray.mojo | 8 +++++++- tests/core/test_setitem_scalar.mojo | 11 ++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/numojo/core/ndarray.mojo b/numojo/core/ndarray.mojo index f7ceeec7..96c0e861 100644 --- a/numojo/core/ndarray.mojo +++ b/numojo/core/ndarray.mojo @@ -2643,7 +2643,13 @@ struct NDArray[dtype: DType = DType.float64]( message=String( "Integer index {} out of bounds for axis {}" " (size {}). Valid range: [-{}, {})." - ).format(idx, i, self.shape[i], self.shape[i], self.shape[i]), + ).format( + idx, + i, + self.shape[i], + self.shape[i], + self.shape[i], + ), location=( "NDArray.__setitem__(*slices: Variant[Slice," " Int], scalar: Scalar)" diff --git a/tests/core/test_setitem_scalar.mojo b/tests/core/test_setitem_scalar.mojo index b81b94ad..ade07b33 100644 --- a/tests/core/test_setitem_scalar.mojo +++ b/tests/core/test_setitem_scalar.mojo @@ -34,7 +34,8 @@ def test_setitem_list_slice_scalar_1d() raises: def test_setitem_list_slice_scalar_2d_submatrix() raises: - """List[Slice] scalar: fills a 2x2 sub-matrix, surrounding values unchanged.""" + """List[Slice] scalar: fills a 2x2 sub-matrix, surrounding values unchanged. + """ var a = nm.arange[nm.i32](0, 16, step=1).reshape(Shape(4, 4)) var sl = List[Slice]() sl.append(Slice(1, 3)) @@ -152,7 +153,9 @@ def test_setitem_variadic_slice_scalar_2d() raises: assert_equal( Int(a.item(i, j)), Int(b.item(i, j)), - String("*Slice and List[Slice] scalar paths diverge at ({},{})").format(i, j), + String( + "*Slice and List[Slice] scalar paths diverge at ({},{})" + ).format(i, j), ) @@ -301,7 +304,9 @@ def test_setitem_mixed_agrees_with_list_slice_scalar() raises: assert_equal( Int(a.item(i, j)), Int(b.item(i, j)), - String("Mixed and List[Slice] scalar paths diverge at ({},{})").format(i, j), + String( + "Mixed and List[Slice] scalar paths diverge at ({},{})" + ).format(i, j), )