Source code for arrayfire.bcast

#######################################################
# 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
########################################################

"""
Function to perform broadcasting operations.
"""

class _bcast(object):
    _flag = False
    def get(self):
        return _bcast._flag

    def set(self, flag):
        _bcast._flag = flag

    def toggle(self):
        _bcast._flag ^= True

_bcast_var = _bcast()

def broadcast(func, *args):
    """
    Function to perform broadcast operations.

    This function can be used directly or as an annotation in the following manner.

    Example
    -------

    Using broadcast as an annotation

    >>> import arrayfire as af
    >>> @af.broadcast
    ... def add(a, b):
    ...     return a + b
    ...
    >>> a = af.randu(2,3)
    >>> b = af.randu(2,1) # b is a different size
    >>> # Trying to add arrays of different sizes raises an exceptions
    >>> c = add(a, b) # This call does not raise an exception because of the annotation
    >>> af.display(a)
    [2 3 1 1]
        0.4107     0.9518     0.4198
        0.8224     0.1794     0.0081

    >>> af.display(b)
    [2 1 1 1]
        0.7269
        0.7104

    >>> af.display(c)
    [2 3 1 1]
        1.1377     1.6787     1.1467
        1.5328     0.8898     0.7185

    Using broadcast as function

    >>> import arrayfire as af
    >>> add = lambda a,b: a + b
    >>> a = af.randu(2,3)
    >>> b = af.randu(2,1) # b is a different size
    >>> # Trying to add arrays of different sizes raises an exceptions
    >>> c = af.broadcast(add, a, b) # This call does not raise an exception
    >>> af.display(a)
    [2 3 1 1]
        0.4107     0.9518     0.4198
        0.8224     0.1794     0.0081

    >>> af.display(b)
    [2 1 1 1]
        0.7269
        0.7104

    >>> af.display(c)
    [2 3 1 1]
        1.1377     1.6787     1.1467
        1.5328     0.8898     0.7185

    """

    def wrapper(*func_args):
        _bcast_var.toggle()
        res = func(*func_args)
        _bcast_var.toggle()
        return res

    if len(args) == 0:
        return wrapper
    else:
        return wrapper(*args)