def test_gufunc(): x = da.random.normal(size=(10, 5), chunks=(2, 5)) def foo(x): return np.mean(x, axis=-1) gufoo = gufunc(foo, signature="(i)->()", output_dtypes=float, vectorize=True) y = gufoo(x) valy = y.compute() assert isinstance(y, Array) assert valy.shape == (10,)
def test_gufunc(): x = da.random.normal(size=(10, 5), chunks=(2, 5)) def foo(x): return np.mean(x, axis=-1) gufoo = gufunc(foo, signature="(i)->()", axis=-1, keepdims=False, output_dtypes=float, vectorize=True) y = gufoo(x) valy = y.compute() assert isinstance(y, Array) assert valy.shape == (10,)