Commits

Taavi Burns committed 78d2ef3

Adds ones_like and zeros_like ufuncs.

Comments (0)

Files changed (4)

pypy/module/micronumpy/__init__.py

         ('logaddexp2', 'logaddexp2'),
         ('real', 'real'),
         ('imag', 'imag'),
+        ('ones_like', 'ones_like'),
+        ('zeros_like', 'zeros_like'),
     ]:
         interpleveldefs[exposed] = "interp_ufuncs.get(space).%s" % impl
 

pypy/module/micronumpy/interp_ufuncs.py

                                        "allow_complex": False}),
             ("logaddexp2", "logaddexp2", 2, {"promote_to_float": True,
                                        "allow_complex": False}),
+            ("ones_like", "ones_like", 1),
+            ("zeros_like", "zeros_like", 1),
         ]:
             self.add_ufunc(space, *ufunc_def)
 

pypy/module/micronumpy/test/test_ufuncs.py

         assert logaddexp2(float('inf'), float('-inf')) == float('inf')
         assert logaddexp2(float('inf'), float('inf')) == float('inf')
 
+    def test_ones_like(self):
+        from numpypy import array, complex, ones_like
 
+        assert ones_like(2) == array(1)
+        assert ones_like(complex(2)) == array(complex(1))
+
+    def test_zeros_like(self):
+        from numpypy import array, complex, zeros_like
+
+        assert zeros_like(2) == array(0)
+        assert zeros_like(complex(2)) == array(complex(0))

pypy/module/micronumpy/types.py

         else:
             return v
 
+    @simple_unary_op
+    def ones_like(self, v):
+        return 1
+
+    @simple_unary_op
+    def zeros_like(self, v):
+        return 0
+
+
 class NonNativePrimitive(Primitive):
     _mixin_ = True
 
         except ValueError:
             return rfloat.NAN, rfloat.NAN
 
+    @complex_unary_op
+    def ones_like(self, v):
+        return 1, 0
+
+    @complex_unary_op
+    def zeros_like(self, v):
+        return 0, 0
+
+
 class Complex64(ComplexFloating, BaseType):
     _attrs_ = ()