MAINT: an incremental simplification of wrapping of unary ufunc

Author: ev-br

autogen/gen_ufuncs_2.py

@@ -85,17 +85,29 @@ def test_{np_name}():
 main_text = header
 test_text = test_header
 
+_impl_list = []
+_ufunc_list = []
+
 for ufunc in dct['ufunc']:
     if ufunc in skip:
         continue
 
     if ufunc.nin == 1:
-        # print(get_signature(ufunc))
+#        print(get_signature(ufunc))
 
         torch_name = torch_names.get(ufunc)
         if torch_name is None:
             torch_name = ufunc.__name__
 
+#        print(ufunc.__name__, ' -- ', torch_name)   
+
+        _impl_stanza = "{np_name} = deco_unary_ufunc(torch.{torch_name})"
+        _impl_stanza = _impl_stanza.format(np_name=ufunc.__name__,
+                                     torch_name=torch_name,)
+        _impl_list.append(_impl_stanza)        
+
+        continue
+
         torch_stanza = stanzas.get(ufunc)
         if torch_stanza is None:
             torch_stanza = default_stanza.format(torch_name=torch_name)
@@ -109,6 +121,11 @@ def test_{np_name}():
 
         _all_list.append(ufunc.__name__)
 
+
+print("\n".join(_impl_list))
+print("\n\n-----\n\n")
+
+'''
 main_text += "\n\n__all__ = %s" % _all_list
 
 
@@ -117,7 +134,7 @@ def test_{np_name}():
 
 with open("test_unary_ufuncs.py", "w") as f:
     f.write(test_text)
-
+'''
 
 ###### BINARY UFUNCS ###################################
 
@@ -163,6 +180,7 @@ def test_{np_name}():
 _impl_list = []
 _ufunc_list = []
 
+
 for ufunc in dct['ufunc']:
 
     if ufunc in skip:
@@ -190,3 +208,6 @@ def test_{np_name}():
 print("\n\n")
 print("\n".join(_ufunc_list))
 
+
+
+

torch_np/_binary_ufuncs.py

@@ -24,7 +24,8 @@
 # files, doing it would currently create import cycles.
 #
 
-
+# TODO: deduplicate with _unary_ufuncs/deco_unary_ufunc_from_impl,
+# _ndarray/asarray_replacer, and _wrapper/concatenate et al
 def deco_ufunc_from_impl(impl_func):
     @functools.wraps(impl_func)
     def wrapped(x1, x2, *args, **kwds):
@@ -75,3 +76,4 @@ def wrapped(x1, x2, *args, **kwds):
 right_shift = deco_ufunc_from_impl(_ufunc_impl.right_shift)
 subtract = deco_ufunc_from_impl(_ufunc_impl.subtract)
 divide = deco_ufunc_from_impl(_ufunc_impl.divide)
+

torch_np/_ufunc_impl.py

@@ -19,15 +19,38 @@ def wrapped(x1, x2, /, out=None, *, where=True,
         # XXX: dtype=... parameter is silently ignored
 
         arrays = (x1, x2)
-        x1_tensor, x2_tensor = _helpers.cast_and_broadcast(arrays, out, casting)
+        tensors = _helpers.cast_and_broadcast(arrays, out, casting)
 
-        result = torch_func(x1_tensor, x2_tensor)
+        result = torch_func(*tensors)
 
         return _helpers.result_or_out(result, out)
     return wrapped
 
 
-# the list is autogenerated, cf autogen/gen_ufunc_2.py
+def deco_unary_ufunc(torch_func):
+    # TODO: deduplicate with `deco_binary_ufunc` above. Need to figure out the
+    #       effect of the `dtype` parameter, does it differ between unary and binary ufuncs. 
+    def wrapped(x1, /, out=None, *, where=True,
+                casting='same_kind', order='K', dtype=None, subok=False, **kwds):
+        _util.subok_not_ok(subok=subok)
+        if order != 'K' or not where:
+            raise NotImplementedError
+
+        # XXX: dtype=... parameter is silently ignored
+
+        arrays = (x1, )
+        tensors  = _helpers.cast_and_broadcast(arrays, out, casting)
+
+        result = torch_func(*tensors)
+
+        return _helpers.result_or_out(result, out)
+    return wrapped
+
+
+
+
+# binary ufuncs: the list is autogenerated, cf autogen/gen_ufunc_2.py
+# And edited manually! np.equal  <--> torch.eq, not torch.equal
 add = deco_binary_ufunc(torch.add)
 arctan2 = deco_binary_ufunc(torch.arctan2)
 bitwise_and = deco_binary_ufunc(torch.bitwise_and)
@@ -69,3 +92,59 @@ def wrapped(x1, x2, /, out=None, *, where=True,
 subtract = deco_binary_ufunc(torch.subtract)
 divide = deco_binary_ufunc(torch.divide)
 
+
+
+# unary ufuncs: the list is autogenerated, cf autogen/gen_ufunc_2.py
+absolute = deco_unary_ufunc(torch.absolute)
+#absolute = deco_unary_ufunc(torch.absolute)
+arccos = deco_unary_ufunc(torch.arccos)
+arccosh = deco_unary_ufunc(torch.arccosh)
+arcsin = deco_unary_ufunc(torch.arcsin)
+arcsinh = deco_unary_ufunc(torch.arcsinh)
+arctan = deco_unary_ufunc(torch.arctan)
+arctanh = deco_unary_ufunc(torch.arctanh)
+ceil = deco_unary_ufunc(torch.ceil)
+conjugate = deco_unary_ufunc(torch.conj_physical)
+#conjugate = deco_unary_ufunc(torch.conj_physical)
+cos = deco_unary_ufunc(torch.cos)
+cosh = deco_unary_ufunc(torch.cosh)
+deg2rad = deco_unary_ufunc(torch.deg2rad)
+degrees = deco_unary_ufunc(torch.rad2deg)
+exp = deco_unary_ufunc(torch.exp)
+exp2 = deco_unary_ufunc(torch.exp2)
+expm1 = deco_unary_ufunc(torch.expm1)
+fabs = deco_unary_ufunc(torch.absolute)
+floor = deco_unary_ufunc(torch.floor)
+isfinite = deco_unary_ufunc(torch.isfinite)
+isinf = deco_unary_ufunc(torch.isinf)
+isnan = deco_unary_ufunc(torch.isnan)
+log = deco_unary_ufunc(torch.log)
+log10 = deco_unary_ufunc(torch.log10)
+log1p = deco_unary_ufunc(torch.log1p)
+log2 = deco_unary_ufunc(torch.log2)
+logical_not = deco_unary_ufunc(torch.logical_not)
+negative = deco_unary_ufunc(torch.negative)
+rad2deg = deco_unary_ufunc(torch.rad2deg)
+radians = deco_unary_ufunc(torch.deg2rad)
+reciprocal = deco_unary_ufunc(torch.reciprocal)
+rint = deco_unary_ufunc(torch.round)
+sign = deco_unary_ufunc(torch.sign)
+signbit = deco_unary_ufunc(torch.signbit)
+sin = deco_unary_ufunc(torch.sin)
+sinh = deco_unary_ufunc(torch.sinh)
+sqrt = deco_unary_ufunc(torch.sqrt)
+square = deco_unary_ufunc(torch.square)
+tan = deco_unary_ufunc(torch.tan)
+tanh = deco_unary_ufunc(torch.tanh)
+trunc = deco_unary_ufunc(torch.trunc)
+
+# special cases: torch does not export these names
+def _cbrt(x):
+    return torch.pow(x, 1/3)
+
+def _positive(x):
+    return +x
+
+cbrt = deco_unary_ufunc(_cbrt)
+positive = deco_unary_ufunc(_positive)
+