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)