-
-
Notifications
You must be signed in to change notification settings - Fork 9.6k
ufunc and binop overrides and behavior suggestions
Marten van Kerkwijk edited this page May 28, 2015
·
31 revisions
This page lists several suggestions on how __numpy_ufunc__, __array_priority__, ndarray subclassing, and Python binary operations should work together in the future.
The discussion is ongoing in gh-5844. As sometimes code speaks more than words, and it is clumsy to revise an evolving suggestion in Github discussion thread, the latest suggestions can be kept here.
Let us assign a name and a revision number for each suggestion. Increment revision number on each edit.
class MyNdarray(object):
_known_classes = (np.ndarray,)
def _can_handle(self, other):
if isinstance(other, _known_classes + (MyNdarray,)):
return True
if hasattr(other, "__numpy_ufunc__"):
return False
return True
def __add__(self, other):
return np.add(self, other)
def __radd__(self, other):
return np.add(other, self)
def __iadd__(self, other):
return np.add(self, other, out=self)
# ... repeat for other binops ...
def __numpy_ufunc__(self, ufunc, method, args, kwargs):
if method == "__call__":
return self._numpy_ufunc_call(ufunc, *args, **kwargs)
elif method == "reduce":
return self._numpy_ufunc_reduce(ufunc, *args, **kwargs)
else:
raise NotImplementedError("it's just an example ok")
def _numpy_ufunc_call(self, ufunc, *args, out=None, **kwargs):
new_args = []
for arg in args:
if not self._can_handle(arg):
return NotImplemented
new_args.append(np.asarray(arg))
args = new_args
if out is not None:
if not isinstance(out, tuple):
out = (out,)
new_out = []
for arg in out:
if not self._can_handle(arg):
return NotImplemented
new_out.append(np.asarray(arg))
out = tuple(new_out)
return ufunc(*args, out=out, **kwargs)
def _numpy_ufunc_reduce(self, a, out=None, **kwargs):
if not self._can_handle(a):
return NotImplemented
else:
a = np.asarray(a)
if out is not None:
if not self._can_handle(out):
return NotImplemented
else:
out = np.asarray(out)
return ufunc.reduce(a, out=out, **kwargs)
class MyNdarray(object):
__numpy_binop_override__ = True
_known_classes = (np.ndarray,)
def _can_handle(self, other):
if isinstance(other, self._known_classes + (MyNdarray,)):
return True
if getattr(other, "__numpy_binop_override__", False):
return False
# potential addition/alternative: consider presence of
# __numpy_ufunc__ as an opt-out
return True
def __add__(self, other):
return self.__numpy_ufunc__(self.add, "__call__", (self, other), {})
def __radd__(self, other):
return self.__numpy_ufunc__(self.add, "__call__", (other, self), {})
def __iadd__(self, other):
# IF we want to never give `other` the opportunity to do an out-of-place
# operation:
return np.add(self, other, out=self)
# OTHERWISE:
return self.__numpy_ufunc__(self.add, "__call__", (self, other), out=self)
# ... repeat for every binop ...
# Implementation of __numpy_ufunc__ is the same as in ufunc.1 above
def __numpy_ufunc__(self, ufunc, method, args, kwargs):
if method == "__call__":
return self._numpy_ufunc_call(ufunc, *args, **kwargs)
elif method == "reduce":
return self._numpy_ufunc_reduce(ufunc, *args, **kwargs)
else:
raise NotImplementedError("it's just an example ok")
def _numpy_ufunc_call(self, ufunc, *args, out=None, **kwargs):
new_args = []
for arg in args:
if not self._can_handle(arg):
return NotImplemented
new_args.append(np.asarray(arg))
args = new_args
if out is not None:
if not isinstance(out, tuple):
out = (out,)
new_out = []
for arg in out:
if not self._can_handle(arg):
return NotImplemented
new_out.append(np.asarray(arg))
out = tuple(new_out)
return ufunc(*args, out=out, **kwargs)
def _numpy_ufunc_reduce(self, a, out=None, **kwargs):
if not self._can_handle(a):
return NotImplemented
else:
a = np.asarray(a)
if out is not None:
if not self._can_handle(out):
return NotImplemented
else:
out = np.asarray(out)
return ufunc.reduce(a, out=out, **kwargs)
Here, I first try to define how ndarray itself behaves, and then give an example of a container subclass as well as a different array-like (the latter is nearly identical to opt-out above).
class ndarray():
def _can_handle(self, other):
if isinstance(self, other): # we're a subclass (or the same), so of course we can handle it.
return True
# Strict subclasses of us and objects defining `__numpy_ufunc__` get treated identically.
if isinstance(other, type(self)) or hasattr(other, '__numpy_ufunc__'):
return False
# backwards compatibility
if getattr(other, '__array_priority__', 0) > 0:
return False
return True
def _can_handle_as_output(self, other):
return isinstance(other, type(self)) and self._can_handle(other)
def __add__(self, other):
return self.__numpy_ufunc__(np.add, "__call__", 0, (self, other), {})
def __radd__(self, other):
return self.__numpy_ufunc__(np.add, "__call__", 1, (other, self), {})
def __iadd__(self, other):
return self.__numpy_ufunc__(np.add, "__call__", 0, (self, other), out=self)
# ... repeat for every binop ...
def __numpy_ufunc__(self, ufunc, method, i, args, **kwargs):
new_args = []
for arg in args:
if not self._can_handle(arg):
return NotImplemented
new_args.append(np.asarray(arg))
args = new_args
if kwargs.pop('out', None) is not None:
if not isinstance(out, tuple):
out = (out,)
new_out = []
for arg in out:
if not self._can_handle_as_output(arg):
return NotImplemented
new_out.append(np.asarray(arg))
kwargs['out'] = tuple(new_out)
# All arguments are now guaranteed to be ndarray, so could call into
# function beyond any checks for `__numpy_ufunc__`, etc.
kwargs['ndarray_only'] = True
return getattr(ufunc, method)(*args, **kwargs)
class ufunc():
def __init__(self, ndarray_only_routine):
# the ndarray_only_routine is the basic code that can only handle ndarray, i.e.,
# does not do any coercing any more.
self.ndarray_only_routine = ndarray_only_routine
def __call__(self, *args, **kwargs):
return self.execute('__call__', *args, **kwargs)
def execute(self, method, *args, **kwargs):
#
<some code that ensures that args only contains inputs and kwargs everything else>
#
if kwargs.pop('ndarray_only', False):
return self.ndarray_only_routine(method, *args, **kwargs)
# we may have non-ndarray instances; try __numpy_ufunc__ until something succeeds
out = kwargs.get('out', ())
if not isinstance(out, tuple):
out = (out,)
result = None
for i, arg in enumerate(args + out):
if hasattr(arg, '__numpy_ufunc__'):
result = arg.__numpy_ufunc__(self, method, i, args, **kwargs)
if result is not NotImplemented:
return result
if result is None:
# None of the arguments were array-like (e.g., np.add(1., 1.); so nothing has been tried
# yet. Use ndarray to see if the arguments can be converted.
result = ndarray.__numpy_ufunc__(self, method, len(args), *args, **kwargs)
if result is not NotImplemented:
return result
raise TypeError
add = ufunc(ndarray_only_add)
class MaskedArray(ndarray):
# A container class that defers everything except masks to ndarray
def __numpy_ufunc__(self, ufunc, method, i, args, **kwargs):
new_args = []
masks = []
for arg in args:
if hasattr(arg, 'mask'): # Duck-typing
masks.append(arg.mask)
new_args.append(arg.data)
else:
masks.append(None)
new_args.append(arg)
args = new_args
if kwargs.pop('out', None) is not None:
if not isinstance(out, tuple):
out = (out,)
new_out = []
for arg in out:
if type(out) is type(self):
new_out.append(arg.data)
else:
new_out.append(arg)
kwargs['out'] = tuple(new_out)
# All arguments are now guaranteed not to be type(self) any more, so we call ndarray
# `__numpy_ufunc__` to deal with possible other types, if needed.
result = super(MaskedArray, self).__numpy_ufunc__(ufunc, method, i, args, **kwargs)
# ignore multiple outputs here for this example's sake.
if result is not out:
out = result.view(type(self))
out.mask = self.combine_masks(masks)
return out
class MyNdarray(object):
_known_classes = (np.ndarray,)
def _can_handle(self, other):
if isinstance(other, self._known_classes + (self.__class__,)):
return True
if hasattr(other, "__numpy_ufunc__"):
return False
return True
# These methods are defined just like for `ndarray`
def __add__(self, other):
return self.__numpy_ufunc__(np.add, "__call__", (self, other), {})
def __radd__(self, other):
return self.__numpy_ufunc__(np.add, "__call__", (other, self), {})
def __iadd__(self, other):
return self.__numpy_ufunc__(np.add, "__call__", (self, other), out=self)
# ... repeat for every binop ...
# Implementation of __numpy_ufunc__ is the same as in ufunc.1 above. This non-container
# class insists it has to know how to deal with other classes rather than change itself into
# an ndarray and let other classes try.
def __numpy_ufunc__(self, ufunc, method, args, kwargs):
if method == "__call__":
return self._numpy_ufunc_call(ufunc, *args, **kwargs)
elif method == "reduce":
return self._numpy_ufunc_reduce(ufunc, *args, **kwargs)
else:
raise NotImplementedError("it's just an example ok")
def _numpy_ufunc_call(self, ufunc, *args, out=None, **kwargs):
new_args = []
for arg in args:
if not self._can_handle(arg):
return NotImplemented
new_args.append(np.asarray(arg))
args = new_args
if out is not None:
if not isinstance(out, tuple):
out = (out,)
new_out = []
for arg in out:
if not self._can_handle(arg):
return NotImplemented
new_out.append(np.asarray(arg))
out = tuple(new_out)
result = ufunc(*args, out=out, **kwargs)
<possible stuff to turn result into type(self)>
return result
def _numpy_ufunc_reduce(self, a, out=None, **kwargs):
if not self._can_handle(a):
return NotImplemented
else:
a = np.asarray(a)
if out is not None:
if not self._can_handle(out):
return NotImplemented
else:
out = np.asarray(out)
return ufunc.reduce(a, out=out, **kwargs)