Skip to content

Commit 6dec6a9

Browse files
[NDArray] Add __setitem__ overload for scalar (#355)
This PR implements the following methods to support `arr[0:1, 3:4] = 10.0`, ```mojo def __setitem__( mut self, *slices: Variant[Slice, Int], scalar: Scalar[Self.dtype] ) raises: def __setitem__( mut self, *slices: Slice, scalar: Scalar[Self.dtype] ) raises: def _setitem_slice_scalar( mut self, slices: List[Slice], val: Scalar[Self.dtype] ) raises ``` This let's the user do, ```mojo def main() raises: var arr = nm.linspace[nm.f32](0, 9, 10).reshape(Shape(2, 5)) print(arr) arr.__setitem__(0, Slice(0, 3, 1), scalar=Scalar[nm.f32](42.0)) print(arr) ``` ## Note: - Unfortunately overload resolution in Mojo isn't fixed yet. So we have to currently call `__setitem__` explicitly with scalar keyword :(. Hopefully once that's fixed, we should be able to do `arr[0:2, 2] = 10.0`.
1 parent 409e7cc commit 6dec6a9

2 files changed

Lines changed: 498 additions & 8 deletions

File tree

numojo/core/ndarray.mojo

Lines changed: 184 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2504,6 +2504,177 @@ struct NDArray[dtype: DType = DType.float64](
25042504

25052505
self.__setitem__(slices=slice_list, val=val)
25062506

2507+
def _setitem_slice_scalar(
2508+
mut self, slices: List[Slice], val: Scalar[Self.dtype]
2509+
) raises:
2510+
"""Internal backend: fill every element in a slice region with a scalar.
2511+
2512+
Called by the user-facing `__setitem__(*Slice, scalar=)` and
2513+
`__setitem__(*Variant[Slice,Int], scalar=)` overloads after they
2514+
normalise their arguments into a `List[Slice]`.
2515+
2516+
Args:
2517+
slices: One `Slice` per array dimension (trailing dims already
2518+
padded to full range by the caller).
2519+
val: The scalar value to write to every selected position.
2520+
"""
2521+
var n_slices: Int = len(slices)
2522+
var slice_list: List[InternalSlice] = self._adjust_slice(slices)
2523+
2524+
for i in range(n_slices, self.ndim):
2525+
slice_list.append(InternalSlice(0, self.shape[i], 1))
2526+
2527+
# Use a single flat counter + coordinate reconstruction to stay
2528+
# allocation-free and layout-agnostic (works for C, F, strided).
2529+
var total: Int = 1
2530+
var region_shape = List[Int](capacity=self.ndim)
2531+
for i in range(self.ndim):
2532+
var slen: Int
2533+
var s = slice_list[i]
2534+
if s.step > 0:
2535+
slen = max((s.end - s.start + s.step - 1) // s.step, 0)
2536+
else:
2537+
slen = max((s.start - s.end - s.step - 1) // (-s.step), 0)
2538+
region_shape.append(slen)
2539+
total *= slen
2540+
2541+
var coords = List[Int](capacity=self.ndim)
2542+
for _ in range(self.ndim):
2543+
coords.append(0)
2544+
2545+
for _ in range(total):
2546+
var buf_idx: Int = self.offset
2547+
for d in range(self.ndim):
2548+
buf_idx += (
2549+
slice_list[d].start + coords[d] * slice_list[d].step
2550+
) * self.strides[d]
2551+
self._buf.ptr.store(buf_idx, val)
2552+
2553+
for d in range(self.ndim - 1, -1, -1):
2554+
coords[d] += 1
2555+
if coords[d] < region_shape[d]:
2556+
break
2557+
coords[d] = 0
2558+
2559+
def __setitem__(
2560+
mut self, *slices: Slice, scalar: Scalar[Self.dtype]
2561+
) raises:
2562+
"""Sets all elements in the slice region to a scalar value.
2563+
2564+
Delegates to `_setitem_slice_scalar` after packing slices into a list.
2565+
2566+
Note: the trailing keyword is named `scalar` (not `val`) to avoid
2567+
ambiguity with the `*slices: Slice, val: Self` overload during Mojo
2568+
overload resolution.
2569+
2570+
Args:
2571+
slices: Variadic slices, one per dimension (trailing dims default
2572+
to the full range).
2573+
scalar: The scalar value to broadcast into every selected position.
2574+
2575+
Examples:
2576+
```mojo
2577+
var a = nm.arange[nm.i32](16).reshape(nm.Shape(4, 4))
2578+
a.__setitem__(Slice(1,3), Slice(1,3), scalar=99)
2579+
```
2580+
"""
2581+
var slice_list = List[Slice](capacity=slices.__len__())
2582+
for i in range(slices.__len__()):
2583+
slice_list.append(slices[i])
2584+
self._setitem_slice_scalar(slice_list, scalar)
2585+
2586+
def __setitem__(
2587+
mut self, *slices: Variant[Slice, Int], scalar: Scalar[Self.dtype]
2588+
) raises:
2589+
"""Sets elements selected by mixed integer/slice indices to a scalar.
2590+
2591+
Handles two cases:
2592+
- All entries are integers (full coordinate) → direct single-element
2593+
write via `_setitem`, no allocation.
2594+
- Mixed or slice-only → normalise integers to unit-length slices and
2595+
delegate to `__setitem__(List[Slice], Scalar)` for the region fill.
2596+
2597+
Note: the trailing keyword is named `scalar` (not `val`) to avoid
2598+
ambiguity with the `*slices: Variant[Slice, Int], val: Self` overload
2599+
during Mojo overload resolution.
2600+
2601+
Args:
2602+
slices: Variadic mix of `Slice` and `Int` index entries.
2603+
scalar: The scalar value to write.
2604+
2605+
Examples:
2606+
```mojo
2607+
var a = nm.arange[nm.i32](16).reshape(nm.Shape(4, 4))
2608+
a.__setitem__(1, 2, scalar=99) # single element
2609+
a.__setitem__(1, Slice(2,4), scalar=0) # one row, partial column
2610+
a.__setitem__(Slice(1,3), Slice(2,4), scalar=7) # sub-matrix
2611+
```
2612+
"""
2613+
var n = slices.__len__()
2614+
if n > self.ndim:
2615+
raise Error(
2616+
NumojoError(
2617+
category="index",
2618+
message=String(
2619+
"Too many indices: received {} but array has only {}"
2620+
" dimensions."
2621+
).format(n, self.ndim),
2622+
location=(
2623+
"NDArray.__setitem__(*slices: Variant[Slice, Int],"
2624+
" scalar: Scalar)"
2625+
),
2626+
)
2627+
)
2628+
2629+
var slice_list = List[Slice](capacity=self.ndim)
2630+
var count_int: Int = 0
2631+
var coords = List[Int]()
2632+
2633+
# Track which array dimension each entry maps to (Int and Slice both
2634+
# consume one dimension; the loop index equals the array dim here
2635+
# because Variant[Slice,Int] carries no NewAxis).
2636+
for i in range(n):
2637+
if slices[i].isa[Int]():
2638+
var idx = slices[i][Int]
2639+
if idx >= self.shape[i] or idx < -self.shape[i]:
2640+
raise Error(
2641+
NumojoError(
2642+
category="index",
2643+
message=String(
2644+
"Integer index {} out of bounds for axis {}"
2645+
" (size {}). Valid range: [-{}, {})."
2646+
).format(
2647+
idx,
2648+
i,
2649+
self.shape[i],
2650+
self.shape[i],
2651+
self.shape[i],
2652+
),
2653+
location=(
2654+
"NDArray.__setitem__(*slices: Variant[Slice,"
2655+
" Int], scalar: Scalar)"
2656+
),
2657+
)
2658+
)
2659+
if idx < 0:
2660+
idx += self.shape[i]
2661+
count_int += 1
2662+
coords.append(idx)
2663+
slice_list.append(Slice(idx, idx + 1, 1))
2664+
else:
2665+
slice_list.append(slices[i][Slice])
2666+
2667+
# Fast path: every dimension was given an integer → single element.
2668+
if count_int == self.ndim:
2669+
self.itemset(coords.copy(), scalar)
2670+
return
2671+
2672+
# Pad trailing dimensions.
2673+
for i in range(n, self.ndim):
2674+
slice_list.append(Slice(0, self.shape[i], 1))
2675+
2676+
self._setitem_slice_scalar(slice_list, scalar)
2677+
25072678
def __setitem__(
25082679
mut self, index: NDArray[DType.int], val: NDArray[Self.dtype]
25092680
) raises:
@@ -3700,20 +3871,25 @@ struct NDArray[dtype: DType = DType.float64](
37003871
)
37013872
else:
37023873
try:
3874+
var order: String
3875+
if self.is_c_contiguous():
3876+
order = "C"
3877+
elif self.is_f_contiguous():
3878+
order = "F"
3879+
else:
3880+
order = "non-contiguous"
37033881
writer.write(
37043882
self._array_to_string(0, 0)
37053883
+ "\n"
37063884
+ String(self.ndim)
3707-
+ "D-array Shape"
3708-
+ String(self.shape)
3709-
+ " Strides"
3710-
+ String(self.strides)
3885+
+ "D-array Shape: "
3886+
+ self.shape.__str__()
3887+
+ " Strides: "
3888+
+ self.strides.__str__()
37113889
+ " DType: "
37123890
+ _concise_dtype_str(self.dtype)
3713-
+ " C-cont: "
3714-
+ String(self.is_c_contiguous())
3715-
+ " F-cont: "
3716-
+ String(self.is_f_contiguous())
3891+
+ " order: "
3892+
+ order
37173893
+ " own data: "
37183894
+ String(self.flags.OWNDATA)
37193895
)

0 commit comments

Comments
 (0)