Skip to content

Commit 4441eb5

Browse files
committed
FEAT: Added the ability to make broadcast a decorator
1 parent 870ee50 commit 4441eb5

File tree

2 files changed

+18
-5
lines changed

2 files changed

+18
-5
lines changed

arrayfire/broadcast.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,14 @@ def toggle():
2020
bcast._flag ^= True
2121

2222
def broadcast(func, *args):
23-
bcast.toggle()
24-
res = func(*args)
25-
bcast.toggle()
26-
return res
23+
24+
def wrapper(*func_args):
25+
bcast.toggle()
26+
res = func(*func_args)
27+
bcast.toggle()
28+
return res
29+
30+
if len(args) == 0:
31+
return wrapper
32+
else:
33+
return wrapper(*args)

tests/simple_arith.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -190,8 +190,14 @@
190190
af.display(af.isnan(a/a))
191191

192192
a = af.randu(5, 1)
193-
b = af.randu(5, 5)
193+
b = af.randu(1, 5)
194194
c = af.broadcast(lambda x,y: x+y, a, b)
195195
af.display(a)
196196
af.display(b)
197197
af.display(c)
198+
199+
@af.broadcast
200+
def test_add(aa, bb):
201+
return aa + bb
202+
203+
af.display(test_add(a, b))

0 commit comments

Comments
 (0)