lcompilers / lpython

Python compiler
https://lpython.org/
Other
1.51k stars 164 forks source link

How to implement sqrt(x) for x an array #242

Open certik opened 2 years ago

certik commented 2 years ago

For a 1D array:

from ltypes import f64, TypeVar
n = TypeVar('n')

@overload
def sqrt(x: f64[n]) -> f64[n]:
    r: f64[n] = empty(size(x))
    i: i32
    for i in range(size(x)):
        r[i] = sqrt(x[i])
    return r

For a 2d array:

from ltypes import f64, TypeVar
n = TypeVar('n')
m = TypeVar('m')

@overload
def sqrt(x: f64[n,m]) -> f64[n,m]:
    r: f64[n,m] = empty((size(x,1),size(x,2)))  # or empty(shape(x))
    i: i32
    j: i32
    for i in range(size(x,1)):
        for j in range(size(x,2)):
            r[i,j] = sqrt(x[i,j])
    return r
certik commented 2 years ago

For now, let's pass n as an argument.