[pypy-commit] pypy ufuncapi: add kw support to generic ufuncs, lay groundwork for kw support in all ufuncs

mattip noreply at buildbot.pypy.org
Fri Oct 10 16:13:42 CEST 2014


Author: mattip <matti.picus at gmail.com>
Branch: ufuncapi
Changeset: r73881:8c3df8a4ea5f
Date: 2014-10-08 17:53 +0300
http://bitbucket.org/pypy/pypy/changeset/8c3df8a4ea5f/

Log:	add kw support to generic ufuncs, lay groundwork for kw support in
	all ufuncs

diff --git a/pypy/module/micronumpy/boxes.py b/pypy/module/micronumpy/boxes.py
--- a/pypy/module/micronumpy/boxes.py
+++ b/pypy/module/micronumpy/boxes.py
@@ -30,6 +30,7 @@
 long_double_size = 8
 
 
+
 def new_dtype_getter(num):
     @specialize.memo()
     def _get_dtype(space):
@@ -200,25 +201,30 @@
     def descr_nonzero(self, space):
         return space.wrap(self.get_dtype(space).itemtype.bool(self))
 
+    # TODO: support all kwargs in ufuncs like numpy ufunc_object.c
+    sig = None
+    cast = None
+    extobj = None
+
     def _unaryop_impl(ufunc_name):
         def impl(self, space, w_out=None):
             from pypy.module.micronumpy import ufuncs
             return getattr(ufuncs.get(space), ufunc_name).call(
-                space, [self, w_out])
+                space, [self, w_out], self.sig, self.cast, self.extobj)
         return func_with_new_name(impl, "unaryop_%s_impl" % ufunc_name)
 
     def _binop_impl(ufunc_name):
         def impl(self, space, w_other, w_out=None):
             from pypy.module.micronumpy import ufuncs
             return getattr(ufuncs.get(space), ufunc_name).call(
-                space, [self, w_other, w_out])
+                space, [self, w_other, w_out], self.sig, self.cast, self.extobj)
         return func_with_new_name(impl, "binop_%s_impl" % ufunc_name)
 
     def _binop_right_impl(ufunc_name):
         def impl(self, space, w_other, w_out=None):
             from pypy.module.micronumpy import ufuncs
             return getattr(ufuncs.get(space), ufunc_name).call(
-                space, [w_other, self, w_out])
+                space, [w_other, self, w_out], self.sig, self.cast, self.extobj)
         return func_with_new_name(impl, "binop_right_%s_impl" % ufunc_name)
 
     descr_add = _binop_impl("add")
diff --git a/pypy/module/micronumpy/compile.py b/pypy/module/micronumpy/compile.py
--- a/pypy/module/micronumpy/compile.py
+++ b/pypy/module/micronumpy/compile.py
@@ -559,10 +559,10 @@
                 w_res = arr.descr_all(interp.space)
             elif self.name == "unegative":
                 neg = ufuncs.get(interp.space).negative
-                w_res = neg.call(interp.space, [arr])
+                w_res = neg.call(interp.space, [arr], None, None, None)
             elif self.name == "cos":
                 cos = ufuncs.get(interp.space).cos
-                w_res = cos.call(interp.space, [arr])
+                w_res = cos.call(interp.space, [arr], None, None, None)
             elif self.name == "flat":
                 w_res = arr.descr_get_flatiter(interp.space)
             elif self.name == "argsort":
diff --git a/pypy/module/micronumpy/ndarray.py b/pypy/module/micronumpy/ndarray.py
--- a/pypy/module/micronumpy/ndarray.py
+++ b/pypy/module/micronumpy/ndarray.py
@@ -855,11 +855,16 @@
         return w_ret
 
     # --------------------- operations ----------------------------
+    # TODO: support all kwargs like numpy ufunc_object.c
+    sig = None
+    cast = None
+    extobj = None
+
 
     def _unaryop_impl(ufunc_name):
         def impl(self, space, w_out=None):
             return getattr(ufuncs.get(space), ufunc_name).call(
-                space, [self, w_out])
+                space, [self, w_out], self.sig, self.cast, self.extobj)
         return func_with_new_name(impl, "unaryop_%s_impl" % ufunc_name)
 
     descr_pos = _unaryop_impl("positive")
@@ -880,7 +885,7 @@
     def _binop_impl(ufunc_name):
         def impl(self, space, w_other, w_out=None):
             return getattr(ufuncs.get(space), ufunc_name).call(
-                space, [self, w_other, w_out])
+                space, [self, w_other, w_out], self.sig, self.cast, self.extobj)
         return func_with_new_name(impl, "binop_%s_impl" % ufunc_name)
 
     descr_add = _binop_impl("add")
@@ -924,7 +929,7 @@
         def impl(self, space, w_other):
             w_out = self
             ufunc = getattr(ufuncs.get(space), ufunc_name)
-            return ufunc.call(space, [self, w_other, w_out])
+            return ufunc.call(space, [self, w_other, w_out], self.sig, self.cast, self.extobj)
         return func_with_new_name(impl, "binop_inplace_%s_impl" % ufunc_name)
 
     descr_iadd = _binop_inplace_impl("add")
@@ -945,7 +950,7 @@
         def impl(self, space, w_other, w_out=None):
             w_other = convert_to_array(space, w_other)
             return getattr(ufuncs.get(space), ufunc_name).call(
-                space, [w_other, self, w_out])
+                space, [w_other, self, w_out], self.sig, self.cast, self.extobj)
         return func_with_new_name(impl, "binop_right_%s_impl" % ufunc_name)
 
     descr_radd = _binop_right_impl("add")
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -112,7 +112,7 @@
             assert 'object' in str(e)
             # Use pypy specific extension for out_dtype
             adder_ufunc0 = frompyfunc(adder, 2, 1, dtypes=['match'])
-            adder_ufunc1 = frompyfunc([adder, adder], 2, 1, 
+            adder_ufunc1 = frompyfunc([adder, adder], 2, 1,
                             dtypes=[int, int, int, float, float, float])
             int_func22 = frompyfunc([int, int], 2, 2, signature='(i),(i)->(i),(i)',
                                     dtypes=['match'])
@@ -147,7 +147,7 @@
             for i in range(in_array.size):
                 out_flat[i] = in_flat[i] * 2
         from numpy import frompyfunc, dtype, arange
-        ufunc = frompyfunc([int_times2, double_times2], 1, 1, 
+        ufunc = frompyfunc([int_times2, double_times2], 1, 1,
                             signature='()->()',
                             dtypes=[dtype(int), dtype(int),
                             dtype(float), dtype(float)
@@ -160,6 +160,21 @@
         af2 = ufunc(af)
         assert all(af2 == af * 2)
 
+    def test_ufunc_kwargs(self):
+        from numpy import ufunc, frompyfunc, arange, dtype
+        def adder(a, b):
+            return a+b
+        adder_ufunc = frompyfunc(adder, 2, 1, dtypes=['match'])
+        args = [arange(10), arange(10)]
+        res = adder_ufunc(*args, dtype=int)
+        assert all(res == args[0] + args[1])
+        # extobj support needed for linalg ufuncs
+        res = adder_ufunc(*args, extobj=[8192, 0, None])
+        assert all(res == args[0] + args[1])
+        raises(TypeError, adder_ufunc, *args, blah=True)
+        raises(TypeError, adder_ufunc, *args, extobj=True)
+        raises(RuntimeError, adder_ufunc, *args, sig='(d,d)->(d)', dtype=int)
+
     def test_ufunc_attrs(self):
         from numpy import add, multiply, sin
 
diff --git a/pypy/module/micronumpy/ufuncs.py b/pypy/module/micronumpy/ufuncs.py
--- a/pypy/module/micronumpy/ufuncs.py
+++ b/pypy/module/micronumpy/ufuncs.py
@@ -27,7 +27,6 @@
         assert isinstance(w_npyobj, W_NDimArray)
         return w_npyobj.get_dtype()
 
-
 class W_Ufunc(W_Root):
     _immutable_fields_ = [
         "name", "promote_to_largest", "promote_to_float", "promote_bools", "nin",
@@ -60,37 +59,37 @@
 
     def descr_call(self, space, __args__):
         args_w, kwds_w = __args__.unpack()
-        # it occurs to me that we don't support any datatypes that
-        # require casting, change it later when we do
-        kwds_w.pop('casting', None)
-        w_subok = kwds_w.pop('subok', None)
-        w_out = kwds_w.pop('out', space.w_None)
-        # Setup a default value for out
+        # sig, extobj are used in generic ufuncs
+        w_subok, w_out, sig, casting, extobj = self.parse_kwargs(space, kwds_w)
         if space.is_w(w_out, space.w_None):
             out = None
         else:
             out = w_out
         if (w_subok is not None and space.is_true(w_subok)):
-            raise OperationError(space.w_NotImplementedError,
-                                 space.wrap("parameters unsupported"))
-        if kwds_w or len(args_w) < self.nin:
-            raise OperationError(space.w_ValueError,
-                space.wrap("invalid number of arguments")
-            )
+            raise oefmt(space.w_NotImplementedError, "parameter subok unsupported")
+        if kwds_w:
+            # numpy compatible, raise with only the first of maybe many keys
+            kw  = kwds_w.keys()[0]
+            raise oefmt(space.w_TypeError,
+                "'%s' is an invalid keyword to ufunc '%s'", kw, self.name)
+        if len(args_w) < self.nin:
+            raise oefmt(space.w_ValueError, "invalid number of arguments"
+                ", expected %d got %d", len(args_w), self.nin)
         elif (len(args_w) > self.nin and out is not None) or \
              (len(args_w) > self.nin + 1):
-            raise OperationError(space.w_TypeError,
-                space.wrap("invalid number of arguments")
-            )
+            raise oefmt(space.w_TypeError, "invalid number of arguments")
         # Override the default out value, if it has been provided in w_wargs
         if len(args_w) > self.nin:
+            if out:
+                raise oefmt(space.w_ValueError, "cannot specify 'out' as both "
+                    "a positional and keyword argument")
             out = args_w[-1]
         else:
             args_w = args_w + [out]
         if out is not None and not isinstance(out, W_NDimArray):
             raise OperationError(space.w_TypeError, space.wrap(
                                             'output must be an array'))
-        return self.call(space, args_w)
+        return self.call(space, args_w, sig, casting, extobj)
 
     def descr_accumulate(self, space, w_obj, w_axis=None, w_dtype=None, w_out=None):
         if space.is_none(w_axis):
@@ -295,6 +294,22 @@
         raise OperationError(space.w_ValueError, space.wrap(
             "outer product only supported for binary functions"))
 
+    def parse_kwargs(self, space, kwds_w):
+        # we don't support casting, change it when we do
+        casting = kwds_w.pop('casting', None)
+        w_subok = kwds_w.pop('subok', None)
+        w_out = kwds_w.pop('out', space.w_None)
+        sig = None
+        # TODO handle triple of extobj,
+        # see _extract_pyvals in ufunc_object.c
+        extobj_w = kwds_w.pop('extobj', get_extobj(space))
+        if not space.isinstance_w(extobj_w, space.w_list) or space.len_w(extobj_w) != 3:
+            raise oefmt(space.w_TypeError, "'extobj' must be a list of 3 values")
+        return w_subok, w_out, sig, casting, extobj_w
+
+def get_extobj(space):
+        extobj_w = space.newlist([space.wrap(8192), space.wrap(0), space.w_None])
+        return extobj_w
 
 class W_Ufunc1(W_Ufunc):
     _immutable_fields_ = ["func", "bool_result"]
@@ -311,7 +326,7 @@
         self.func = func
         self.bool_result = bool_result
 
-    def call(self, space, args_w):
+    def call(self, space, args_w, sig, casting, extobj):
         w_obj = args_w[0]
         out = None
         if len(args_w) > 1:
@@ -397,7 +412,8 @@
         return False
 
     @jit.unroll_safe
-    def call(self, space, args_w):
+    def call(self, space, args_w, sig, casting, extobj):
+        w_obj = args_w[0]
         if len(args_w) > 2:
             [w_lhs, w_rhs, w_out] = args_w
         else:
@@ -529,9 +545,8 @@
                cumulative=False):
         raise oefmt(space.w_NotImplementedError, 'not implemented yet')
 
-    def call(self, space, args_w):
-        #from pypy.module._cffi_backend import newtype, func as _func
-        out = None
+    def call(self, space, args_w, sig, casting, extobj):
+        w_obj = args_w[0]
         inargs = []
         if len(args_w) < self.nin:
             raise oefmt(space.w_ValueError,
@@ -549,7 +564,9 @@
                     raise oefmt(space.w_TypeError,
                          'output arg %d must be an array, not %s', i+self.nin, str(args_w[i+self.nin]))
                 outargs[i] = out
-        index = self.type_resolver(space, inargs, outargs)
+        if sig is None:
+            sig = space.wrap(self.signature)
+        index = self.type_resolver(space, inargs, outargs, sig)
         outargs = self.alloc_outargs(space, index, inargs, outargs)
         inargs0 = inargs[0]
         outargs0 = outargs[0]
@@ -573,9 +590,38 @@
         return loop.call_many_to_many(space, new_shape, self.funcs[index],
                                      res_dtype, inargs, outargs)
 
-    def type_resolver(self, space, inargs, outargs):
+    def parse_kwargs(self, space, kwargs_w):
+        w_subok, w_out, casting, sig, extobj = \
+                    W_Ufunc.parse_kwargs(self, space, kwargs_w)
+        dtype_w = kwargs_w.pop('dtype', None)
+        if not space.is_w(dtype_w, space.w_None) and not dtype_w is None:
+            if sig:
+                raise oefmt(space.w_RuntimeError,
+                        "cannot specify both 'sig' and 'dtype'")
+            dtype = descriptor.decode_w_dtype(space, dtype_w)
+            sig = space.newtuple([dtype])
+        order = kwargs_w.pop('dtype', None)
+        if not space.is_w(order, space.w_None) and not order is None:
+            raise oefmt(space.w_NotImplementedError, '"order" keyword not implemented')
+        parsed_kw = []
+        for kw in kwargs_w:
+            if kw.startswith('sig'):
+                if sig:
+                    raise oefmt(space.w_RuntimeError,
+                            "cannot specify both 'sig' and 'dtype'")
+                sig = kwargs_w[kw]
+                parsed_kw.append(kw)
+            elif kw.startswith('where'):
+                raise oefmt(space.w_NotImplementedError,
+                            '"where" keyword not implemented')
+                parsed_kw.append(kw)
+        for kw in parsed_kw:
+            kwargs_w.pop(kw)
+        return w_subok, w_out, sig, casting, extobj
+
+    def type_resolver(self, space, inargs, outargs, sig):
         # Find a match for the inargs.dtype in self.dtypes, like
-        # linear_search_type_resolver in numy ufunc_type_resolutions.c
+        # linear_search_type_resolver in numpy ufunc_type_resolutions.c
         inargs0 = inargs[0]
         assert isinstance(inargs0, W_NDimArray)
         for i in range(0, len(self.dtypes), self.nargs):
@@ -601,7 +647,7 @@
                 outargs[i] = W_NDimArray.from_shape(space, temp_shape, dtype, order)
         for i in range(len(outargs)):
             assert isinstance(outargs[i], W_NDimArray)
-        return outargs    
+        return outargs
 
     def prep_call(self, space, index, inargs, outargs):
         # Use the index and signature to determine


More information about the pypy-commit mailing list