@@ -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