def test_prim_array_scan():
v = np.ones((2, 3))
def f(a, b):
return a + b
vref = np.cumsum(v, axis=1)
v2 = array_scan(f, 0, v, 1)
assert (v == 1).all()
assert (v2 == vref).all()
def f(x):
def add(x, y):
return x + y
return array_scan(add, 0, x, 1)
def test_array_scan(ary, ax):
def f(a, b):
return a + b
return array_scan(f, 0, ary, ax)