import numpy, ctypes
import numba
import math
from numba.vectorize import GUVectorize

def exp_avg(arr_t, decay_length, out_t):
    decay_factor = math.exp(-1.0 / decay_length)
    sum_w = 0.0
    sum_wx = 0.0
   
    n_t = arr_t.shape[0]
    for t in xrange(n_t):
        sum_w += 1.0
        sum_wx += arr_t[t]
        out_t[t] = sum_wx / sum_w
       
        sum_w *= decay_factor
        sum_wx *= decay_factor

def saxpy(a, X, Y, out):
    for i in range(out.shape[0]):
        out[i] = a * X[i] + Y[i]

def test_exp_avg():
    vect = GUVectorize(exp_avg, '(t),()->(t)')
    vect.add('void(f8[:], f8, f8[:])')
    ufunc = vect.build_ufunc()

    arr_t = numpy.arange(10, dtype=numpy.float64)
    decay_length = numpy.float64(123)
    out_t = numpy.empty(10, dtype=numpy.float64)

    ufunc(arr_t, decay_length, out_t)

    expect = numpy.zeros_like(out_t)
    exp_avg(arr_t, decay_length, expect)

    numpy.allclose(expect, out_t)

def test_saxpy():
    vect = GUVectorize(saxpy, '(),(t),(t)->(t)')
    vect.add('void(f4, f4[:], f4[:], f4[:])')
    gusaxpy = vect.build_ufunc()

    
    A = numpy.array(numpy.float32(2))
    X = numpy.arange(10, dtype=numpy.float32).reshape(5,2)
    Y = numpy.arange(10, dtype=numpy.float32).reshape(5,2)

    out = gusaxpy(A, X, Y)

    for j in range(5):
        for i in range(2):
            exp = A * X[j, i] + Y[j, i]
            assert exp == out[j, i], (exp, out[j, i])

    A = numpy.arange(5, dtype=numpy.float32)
    out = gusaxpy(A, X, Y)

    for j in range(5):
        for i in range(2):
            exp = A[j] * X[j, i] + Y[j, i]
            assert exp == out[j, i]


def main():
    test_exp_avg()
    test_saxpy()

if __name__ == '__main__':
    main()

@@ -1618,16 +1618,583 @@
 j, i%5D%0A%0A%0A
+def mycore(A,B):%0A    m, n = A.shape%0A    B%5B0%5D = 0%0A    for i in range(m):%0A        for j in range(n):%0A            B%5B0%5D += A%5Bi,j%5D%0A%0A%0Adef test_scalar_output():%0A    gufunc = GUVectorize(mycore, '(m,n)-%3E()')%0A    gufunc.add('void(float32%5B:,:%5D, float32%5B:%5D)')%0A    gufunc.add('void(float64%5B:,:%5D, float64%5B:%5D)')%0A%0A    myfunc = gufunc.build_ufunc()%0A%0A    A = numpy.arange(2 * 2 * 10).reshape(10, 2, 2)%0A    out = numpy.zeros(10)%0A    myfunc(A, out=out)%0A%0A    expect = numpy.zeros_like(out)%0A    for i in range(A.shape%5B0%5D):%0A        expect%5Bi%5D = A%5Bi%5D.sum()%0A%0A    numpy.allclose(expect, out)%0A%0A
 def main
@@ -2232,16 +2232,41 @@
 _saxpy()
+%0A    test_scalar_output()
 %0A%0Aif __n

test scalar output support