FEAT: Adding broadcast to arrayfire

Author: pavanky

arrayfire/__init__.py

@@ -22,6 +22,7 @@
 from .features   import *
 from .vision     import *
 from .graphics   import *
+from .broadcast  import *
 
 # do not export default modules as part of arrayfire
 del ct
@@ -34,6 +35,7 @@
 del seq
 del index
 del cell
+del bcast
 
 #do not export internal functions
 del binary_func

arrayfire/arith.py

@@ -9,6 +9,7 @@
 
 from .library import *
 from .array import *
+from .broadcast import *
 
 def arith_binary_func(lhs, rhs, c_func):
     out = array()
@@ -20,21 +21,21 @@ def arith_binary_func(lhs, rhs, c_func):
         TypeError("Atleast one input needs to be of type arrayfire.array")
 
     elif (is_left_array and is_right_array):
-        safe_call(c_func(ct.pointer(out.arr), lhs.arr, rhs.arr, False))
+        safe_call(c_func(ct.pointer(out.arr), lhs.arr, rhs.arr, bcast.get()))
 
     elif (is_number(rhs)):
         ldims = dim4_tuple(lhs.dims())
         lty = lhs.type()
         other = array()
         other.arr = constant_array(rhs, ldims[0], ldims[1], ldims[2], ldims[3], lty)
-        safe_call(c_func(ct.pointer(out.arr), lhs.arr, other.arr, False))
+        safe_call(c_func(ct.pointer(out.arr), lhs.arr, other.arr, bcast.get()))
 
     else:
         rdims = dim4_tuple(rhs.dims())
         rty = rhs.type()
         other = array()
         other.arr = constant_array(lhs, rdims[0], rdims[1], rdims[2], rdims[3], rty)
-        safe_call(c_func(ct.pointer(out.arr), lhs.arr, other.arr, False))
+        safe_call(c_func(ct.pointer(out.arr), lhs.arr, other.arr, bcast.get()))
 
     return out
 

arrayfire/array.py

@@ -10,6 +10,7 @@
 import inspect
 from .library import *
 from .util import *
+from .broadcast import *
 
 def create_array(buf, numdims, idims, dtype):
     out_arr = ct.c_longlong(0)
@@ -63,7 +64,7 @@ def binary_func(lhs, rhs, c_func):
     elif not isinstance(rhs, array):
         raise TypeError("Invalid parameter to binary function")
 
-    safe_call(c_func(ct.pointer(out.arr), lhs.arr, other.arr, False))
+    safe_call(c_func(ct.pointer(out.arr), lhs.arr, other.arr, bcast.get()))
 
     return out
 
@@ -79,7 +80,7 @@ def binary_funcr(lhs, rhs, c_func):
     elif not isinstance(lhs, array):
         raise TypeError("Invalid parameter to binary function")
 
-    c_func(ct.pointer(out.arr), other.arr, rhs.arr, False)
+    c_func(ct.pointer(out.arr), other.arr, rhs.arr, bcast.get())
 
     return out
 

arrayfire/broadcast.py

@@ -0,0 +1,26 @@
+#######################################################
+# Copyright (c) 2015, ArrayFire
+# All rights reserved.
+#
+# This file is distributed under 3-clause BSD license.
+# The complete license agreement can be obtained at:
+# http://arrayfire.com/licenses/BSD-3-Clause
+########################################################
+
+
+class bcast(object):
+    _flag = False
+    def get():
+        return bcast._flag
+
+    def set(flag):
+        bcast._flag = flag
+
+    def toggle():
+        bcast._flag ^= True
+
+def broadcast(func, *args):
+    bcast.toggle()
+    res = func(*args)
+    bcast.toggle()
+    return res

tests/simple_arith.py

@@ -188,3 +188,10 @@
 af.display(af.iszero(a))
 af.display(af.isinf(a/b))
 af.display(af.isnan(a/a))
+
+a = af.randu(5, 1)
+b = af.randu(5, 5)
+c = af.broadcast(lambda x,y: x+y, a, b)
+af.display(a)
+af.display(b)
+af.display(c)