def __call__(self, x1, x2, y=None, *args, **kwargs):
# What types of input are allowed?
x1_isdla = isinstance(x1, DenseLocalArray)
x2_isdla = isinstance(x2, DenseLocalArray)
y_isdla = isinstance(y, DenseLocalArray)
assert x1_isdla or isscalar(x1), "Invalid type for binary ufunc"
assert x2_isdla or isscalar(x2), "Invalid type for binary ufunc"
assert y is None or y_isdla
if y is None:
if x1_isdla and x2_isdla:
if not arecompatible(x1, x2):
raise IncompatibleArrayError("Incompatible DistArrays")
return self.func(x1, x2, *args, **kwargs)
elif y_isdla:
if x1_isdla:
if not arecompatible(x1, y):
raise IncompatibleArrayError("Incompatible LocalArrays")
if x2_isdla:
if not arecompatible(x2, y):
raise IncompatibleArrayError("Incompatible LocalArrays")
kwargs.pop('y', None)
self.func(x1, x2, y.local_array, *args, **kwargs)
return y
else:
raise TypeError("Invalid return type for unary ufunc")
def asdist_like(self, other):
"""
Return a version of self that has shape, dist and grid_shape like
other.
"""
if arecompatible(self, other):
return self
else:
raise IncompatibleArrayError("DistArrays have incompatible shape,"
"dist or grid_shape")
