NDArray implementation in mojo

To develop AI function we need a ndarray like class. But mojo does not have this kind of class now. So we need to
implement one. The most nature way is to extend NDBuffer in mojo. But mojo does not support inheritance now. In this
article I will use composite instead of inheritance to extend NDBuffer.
First we need import necessary libs:

# common import
from String import String
from Bool import Bool
from List import VariadicList
from Buffer import NDBuffer
from List import DimList
from DType import DType
from Pointer import DTypePointer
from TargetInfo import dtype_sizeof, dtype_simd_width
from Index import StaticIntTuple

alias nelts = dtype_simd_width[DType.f32]()

Now it is time to extend NDBuffer:

struct NDArray[rank:Int, dims:DimList, dtype:DType]:
    var ndb:NDBuffer[rank, dims, dtype]
    
    fn __init__(inout self, size:Int):
        let data = DTypePointer[dtype].alloc(size)
        self.ndb = NDBuffer[rank, dims, dtype](data)
        
    fn __getitem__(self, idxs:StaticIntTuple[rank]) -> SIMD[dtype,1]:
        return self.ndb.simd_load[1](idxs)
    
    fn __setitem__(self, idxs:StaticIntTuple[rank], val:SIMD[dtype,1]):
        self.ndb.simd_store[1](idxs, val)

To use the NDArray:

alias x_rank = 2
var x = NDArray[x_rank, DimList(2, 3), DType.f32](6)
x[StaticIntTuple[x_rank](1,1)] = 999.9
print(x[StaticIntTuple[x_rank](0,0)], x[StaticIntTuple[x_rank](1,1)])

Discussion

As see the above code you may be ask why not use x[0,0] to get and set the item of NDArray. The reason is that mojo
dose not support named keyword arguments now. So if I define the NDArray like these:

struct ndarray[rank:Int, dims:DimList, dtype:DType]:
    var ndb:NDBuffer[rank, dims, dtype]
    
    fn __init__(inout self, size:Int):
        let data = DTypePointer[dtype].alloc(size)
        self.ndb = NDBuffer[rank, dims, dtype](data)
        
    fn __getitem__(self, *idxs:Int) -> SIMD[dtype,1]:
        return self.ndb.simd_load[1](idxs)
    
    fn __setitem__(self, *idxs:Int, val:SIMD[dtype,1]):
        self.ndb.simd_store[1](idxs, val)

When I access x[0,0] it will report these errors:

error: Expression [23]:16:37: TODO: keyword arguments not supported yet
    fn __setitem__(self, *idxs:Int, val:SIMD[dtype,1]):
                                    ^

error: Expression [23]:17:31: no matching function in call to 'simd_store': 
        self.ndb.simd_store[1](idxs, val)
        ~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~

/.modular/Kernels/mojo/Stdlib/Buffer.mojo:1173:5: candidate not viable: method argument #0 cannot be converted from 'variadic<_"$Int"::_Int>' to 'StaticIntTuple[rank]'
    fn simd_store[
    ^

/.modular/Kernels/mojo/Stdlib/Buffer.mojo:1191:5: candidate not viable: method argument #0 cannot be converted from 'variadic<_"$Int"::_Int>' to 'StaticTuple[rank, Int]'
    fn simd_store[
    ^