Merge pull request #54 from shivasankarka/ndarray

Minor fixes

Author: MadAlex1997

numojo/core/ndarray.mojo

@@ -521,14 +521,14 @@ struct _NDArrayIter[
     var length: Int
 
     fn __init__(
-        inout self, 
-        unsafe_pointer: DTypePointer[dtype], 
+        inout self,
+        unsafe_pointer: DTypePointer[dtype],
         length: Int,
     ):
         self.index = 0 if forward else length
         self.ptr = unsafe_pointer
         self.length = length
-        
+
     fn __iter__(self) -> Self:
         return self
 
@@ -550,12 +550,15 @@ struct _NDArrayIter[
         else:
             return self.index
 
+
 # ===----------------------------------------------------------------------===#
 # NDArray
 # ===----------------------------------------------------------------------===#
 
 
-struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Sized):
+struct NDArray[dtype: DType = DType.float32](
+    Stringable, CollectionElement, Sized
+):
     """The N-dimensional array (NDArray).
 
     The array can be uniquely defined by the following:
@@ -927,7 +930,7 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
         for i in range(index.__len__()):
             if index[i] >= self.ndshape[i]:
                 raise Error("Error: Elements of `index` exceed the array shape")
-        var idx: Int = _get_index(index, self.coefficient)
+        var idx: Int = _get_index(index, self.stride)
         return self.data.load[width=1](idx)
 
     fn __getitem__(self, index: List[Int]) raises -> SIMD[dtype, 1]:
@@ -1007,6 +1010,8 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
         var count: Int = 0
         for i in range(slices.__len__()):
             self._adjust_slice_(slices[i], self.ndshape[i])
+            if slices[i].start >= self.ndshape[i] or slices[i].end > self.ndshape[i]:
+                raise Error("Error: Slice value exceeds the array shape")
             spec.append(slices[i].unsafe_indices())
             if slices[i].unsafe_indices() != 1:
                 ndims += 1
@@ -1024,7 +1029,7 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
         count = 0
         for _ in range(ndims):
             while spec[j] == 1:
-                count+=1
+                count += 1
                 j += 1
             if j >= self.ndim:
                 break
@@ -1041,8 +1046,6 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
         var noffset: Int = 0
         if self.order == "C":
             noffset = 0
-            if ndims == 1:
-                nstrides.append(1)
             for i in range(ndims):
                 var temp_stride: Int = 1
                 for j in range(i + 1, ndims):  # temp
@@ -1281,7 +1284,7 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
         )
 
     fn __reversed__(self) -> _NDArrayIter[dtype, forward=False]:
-        """Iterate backwards over elements of the NDArray, returning 
+        """Iterate backwards over elements of the NDArray, returning
         copied value.
 
         Returns:
@@ -1519,41 +1522,45 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
         # We might need to figure out how we want to handle truthyness before can do this
         alias nelts: Int = simdwidthof[dtype]()
         var result: Bool = True
+
         @parameter
         fn vectorized_all[simd_width: Int](idx: Int) -> None:
-            result = result and allb(self.data.load[width=simd_width](idx) ) 
+            result = result and allb(self.data.load[width=simd_width](idx))
+
         vectorize[vectorized_all, nelts](self.ndshape._size)
-        return result 
+        return result
 
     fn any(self) raises -> Bool:
         # make this a compile time check
         if not (self.dtype == DType.bool or is_inttype(dtype)):
             raise Error("Array elements must be Boolean or Integer.")
         alias nelts: Int = simdwidthof[dtype]()
-        var result: Bool = False 
+        var result: Bool = False
+
         @parameter
         fn vectorized_any[simd_width: Int](idx: Int) -> None:
-            result = result or anyb(self.data.load[width=simd_width](idx) ) 
+            result = result or anyb(self.data.load[width=simd_width](idx))
+
         vectorize[vectorized_any, nelts](self.ndshape._size)
         return result
 
     fn argmax(self) -> Int:
         var result: Int = 0
         var max_val: SIMD[dtype, 1] = self.load[width=1](0)
         for i in range(1, self.ndshape._size):
-            var temp: SIMD[dtype, 1] = self.load[width=1](i) 
-            if  temp > max_val:
-                max_val = temp 
+            var temp: SIMD[dtype, 1] = self.load[width=1](i)
+            if temp > max_val:
+                max_val = temp
                 result = i
         return result
 
     fn argmin(self) -> Int:
         var result: Int = 0
         var min_val: SIMD[dtype, 1] = self.load[width=1](0)
         for i in range(1, self.ndshape._size):
-            var temp: SIMD[dtype, 1] = self.load[width=1](i) 
-            if  temp < min_val:
-                min_val = temp 
+            var temp: SIMD[dtype, 1] = self.load[width=1](i)
+            if temp < min_val:
+                min_val = temp
                 result = i
         return result
 
@@ -1563,13 +1570,16 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
     fn astype[type: DType](inout self) raises -> NDArray[type]:
         # I wonder if we can do this operation inplace instead of allocating memory.
         alias nelts = simdwidthof[dtype]()
-        var narr: NDArray[type] = NDArray[type](self.ndshape, random=False, order=self.order)
+        var narr: NDArray[type] = NDArray[type](
+            self.ndshape, random=False, order=self.order
+        )
         narr.datatype = type
+
         @parameter
         fn vectorized_astype[width: Int](idx: Int) -> None:
             narr.store[width](idx, self.load[width](idx).cast[type]())
 
-        vectorize[vectorized_astype, nelts](self.ndshape._size)    
+        vectorize[vectorized_astype, nelts](self.ndshape._size)
         return narr
 
     # fn clip(self):
@@ -1619,8 +1629,11 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
         #     self.stride = NDArrayStride(shape = self.ndshape, offset=0)
         #     return self
 
-        var res: NDArray[dtype] = NDArray[dtype](self.ndshape._size, random=False)
+        var res: NDArray[dtype] = NDArray[dtype](
+            self.ndshape._size, random=False
+        )
         alias simd_width: Int = simdwidthof[dtype]()
+
         @parameter
         fn vectorized_flatten[simd_width: Int](index: Int) -> None:
             res.data.store[width=simd_width](

numojo/core/ndarray_utils.mojo

@@ -68,9 +68,7 @@ fn _traverse_iterative[
     if depth == ndim.__len__():
         var idx = offset + _get_index(index, coefficients)
         var nidx = _get_index(index, strides)
-        var temp = orig.data.load[width=1](
-            idx
-        )
+        var temp = orig.data.load[width=1](idx)
         if nidx >= narr.ndshape._size:
             raise Error("Invalid index: index out of bound")
         else:

numojo/core/sort.mojo

@@ -20,11 +20,8 @@ TODO:
 # Bubble sort
 # ===------------------------------------------------------------------------===#
 
-fn bubble_sort[
-    dtype: DType
-    ](
-        ndarray: NDArray[dtype]
-        ) raises -> NDArray[dtype]:
+
+fn bubble_sort[dtype: DType](ndarray: NDArray[dtype]) raises -> NDArray[dtype]:
     """
     Bubble sort the NDArray.
     Average complexity: O(n^2) comparisons, O(n^2) swaps.
@@ -44,11 +41,11 @@ fn bubble_sort[
     var length = ndarray.size()
 
     for i in range(length):
-        for j in range(length-i-1):
-            if result.data.load[width=1](j) > result.data.load[width=1](j+1):
+        for j in range(length - i - 1):
+            if result.data.load[width=1](j) > result.data.load[width=1](j + 1):
                 var temp = result.data.load[width=1](j)
-                result.data.store[width=1](j, result.data.load[width=1](j+1))
-                result.data.store[width=1](j+1, temp)
+                result.data.store[width=1](j, result.data.load[width=1](j + 1))
+                result.data.store[width=1](j + 1, temp)
 
     return result
 
@@ -57,13 +54,10 @@ fn bubble_sort[
 # Quick sort
 # ===------------------------------------------------------------------------===#
 
+
 fn _partition(
-    inout ndarray: NDArray, 
-    left: Int, 
-    right: Int, 
-    pivot_index: Int
-    ) raises -> Int:
-    
+    inout ndarray: NDArray, left: Int, right: Int, pivot_index: Int
+) raises -> Int:
     var pivot_value = ndarray[pivot_index]
     ndarray[pivot_index], ndarray[right] = ndarray[right], ndarray[pivot_index]
     var store_index = left
@@ -73,14 +67,13 @@ fn _partition(
             ndarray[store_index], ndarray[i] = ndarray[i], ndarray[store_index]
             store_index = store_index + 1
     ndarray[right], ndarray[store_index] = ndarray[store_index], ndarray[right]
-    
+
     return store_index
 
-fn quick_sort_inplace[dtype: DType](
-    inout ndarray: NDArray[dtype],
-    left: Int,
-    right: Int,
-    ) raises:
+
+fn quick_sort_inplace[
+    dtype: DType
+](inout ndarray: NDArray[dtype], left: Int, right: Int,) raises:
     """
     Quick sort (in-place) the NDArray.
 
@@ -96,12 +89,13 @@ fn quick_sort_inplace[dtype: DType](
     if right > left:
         var pivot_index = left + (right - left) // 2
         var pivot_new_index = _partition(ndarray, left, right, pivot_index)
-        quick_sort_inplace(ndarray, left, pivot_new_index-1)
-        quick_sort_inplace(ndarray, pivot_new_index+1, right)
+        quick_sort_inplace(ndarray, left, pivot_new_index - 1)
+        quick_sort_inplace(ndarray, pivot_new_index + 1, right)
 
-fn quick_sort[dtype: DType](
-    ndarray: NDArray[dtype],
-    ) raises -> NDArray[dtype]:
+
+fn quick_sort[
+    dtype: DType
+](ndarray: NDArray[dtype],) raises -> NDArray[dtype]:
     """
     Quick sort the NDArray.
     Adopt in-place partition.
@@ -123,14 +117,16 @@ fn quick_sort[dtype: DType](
     var result: NDArray[dtype] = ndarray
     var length = ndarray.size()
 
-    quick_sort_inplace(result, 0, length-1)
+    quick_sort_inplace(result, 0, length - 1)
 
     return result
 
+
 # ===------------------------------------------------------------------------===#
 # Binary sort
 # ===------------------------------------------------------------------------===#
 
+
 fn binary_sort[
     in_dtype: DType, out_dtype: DType = DType.float64
 ](array: NDArray[in_dtype]) raises -> NDArray[out_dtype]:
@@ -158,4 +154,4 @@ fn binary_sort[
                 var temp: Scalar[out_dtype] = result[i - 1]
                 result[i - 1] = result[i]
                 result[i] = temp
-    return result
+    return result

numojo/core/utility_funcs.mojo

@@ -27,6 +27,7 @@ fn is_inttype(dtype: DType) -> Bool:
         return True
     return False
 
+
 fn is_floattype[dtype: DType]() -> Bool:
     if (
         dtype == DType.float16

numojo/math/__init__.mojo

@@ -10,5 +10,5 @@ from ._math_funcs import (
     VectorizedParallelized,
     VectorizedUnroll,
     VectorizedVerbose,
-    VectorizedParallelizedNWorkers
+    VectorizedParallelizedNWorkers,
 )