Wrap arctan2 from numpy, elemwise.

theano/scalar/basic.py

@@ -2113,6 +2113,27 @@ class ArcTan(UnaryScalarOp):
 arctan = ArcTan(upgrade_to_float, name='arctan')
 
 
+class ArcTan2(BinaryScalarOp):
+    def impl(self, y, x):
+        return numpy.arctan2(y, x)
+
+    def grad(self, (y, x), (gz,)):
+        if gz.type in complex_types:
+            raise NotImplementedError()
+        if x.type in float_types and y.type in float_types:
+            return [gz * x / (sqr(x) + sqr(y)),
+                    gz * neg(y) / (sqr(x) + sqr(y))]
+        else:
+            return None,
+
+    def c_code(self, node, name, (y, x), (z,), sub):
+        if (node.inputs[0].type in complex_types or
+            node.inputs[1].type in complex_types):
+            raise NotImplementedError('type not supported', type)
+        return "%(z)s = atan2(%(y)s, %(x)s);" % locals()
+arctan2 = ArcTan2(upgrade_to_float, name='arctan2')
+
+
 class Cosh(UnaryScalarOp):
     """
     cosh(x) = (exp(x) + exp(-x)) / 2

theano/tensor/basic.py

@@ -2596,6 +2596,11 @@ def arctan(a):
     """arctangent of a"""
 
 
+@_scal_elemwise_with_nfunc('arctan2', 1, 1)
+def arctan2(a, b):
+    """arctangent of a / b"""
+
+
 @_scal_elemwise_with_nfunc('cosh', 1, 1)
 def cosh(a):
     """hyperbolic cosine of a"""

theano/tensor/inplace.py

@@ -167,6 +167,10 @@ def tan_inplace(a):
 def arctan_inplace(a):
     """arctangent of `a` (inplace on `a`)"""
 
+@_scal_inplace
+def arctan2_inplace(a, b):
+    """arctangent of `a` / `b` (inplace on `a`)"""
+
 @_scal_inplace
 def cosh_inplace(a):
     """hyperbolic cosine of `a` (inplace on `a`)"""

theano/tensor/tests/test_basic.py

@@ -1103,18 +1103,49 @@ TanInplaceTester = makeBroadcastTester(op=inplace.tan_inplace,
                                        grad_rtol=tan_grad_rtol,
                                        inplace=True)
 
-ArcTanTester = makeBroadcastTester(op=tensor.tan,
-                                   expected=numpy.tan,
+ArcTanTester = makeBroadcastTester(op=tensor.arctan,
+                                   expected=numpy.arctan,
                                    good=_good_broadcast_unary_wide,
                                    grad=_grad_broadcast_unary_wide,
                                    grad_rtol=tan_grad_rtol)
-ArcTanInplaceTester = makeBroadcastTester(op=inplace.tan_inplace,
-                                          expected=numpy.tan,
+ArcTanInplaceTester = makeBroadcastTester(op=inplace.arctan_inplace,
+                                          expected=numpy.arctan,
                                           good=_good_broadcast_unary_wide,
                                           grad=_grad_broadcast_unary_wide,
                                           grad_rtol=tan_grad_rtol,
                                           inplace=True)
 
+_good_broadcast_binary_arctan2 = dict(
+    same_shapes=(rand(2, 3), rand(2, 3)),
+    not_same_dimensions=(rand(2, 2), rand(2)),
+    scalar=(rand(2, 3), rand(1, 1)),
+    row=(rand(2, 3), rand(1, 3)),
+    column=(rand(2, 3), rand(2, 1)),
+    integers=(randint(2, 3), randint(2, 3)),
+    dtype_mixup_1=(rand(2, 3), randint(2, 3)),
+    dtype_mixup_2=(randint(2, 3), rand(2, 3)),
+    empty=(numpy.asarray([]), numpy.asarray([1])),
+    )
+
+_grad_broadcast_binary_arctan2 = dict(
+    same_shapes=(rand(2, 3), rand(2, 3)),
+    scalar=(rand(2, 3), rand(1, 1)),
+    row=(rand(2, 3), rand(1, 3)),
+    column=(rand(2, 3), rand(2, 1)),
+    )
+
+ArcTan2Tester = makeBroadcastTester(op=tensor.arctan2,
+                                    expected=numpy.arctan2,
+                                    good=_good_broadcast_binary_arctan2,
+                                    grad=_grad_broadcast_binary_arctan2,
+                                    grad_rtol=tan_grad_rtol)
+ArcTan2InplaceTester = makeBroadcastTester(op=inplace.arctan2_inplace,
+                                           expected=numpy.arctan2,
+                                           good=_good_broadcast_binary_arctan2,
+                                           grad=_grad_broadcast_binary_arctan2,
+                                           grad_rtol=tan_grad_rtol,
+                                           inplace=True)
+
 CoshTester = makeBroadcastTester(op=tensor.cosh,
                                  expected=numpy.cosh,
                                  good=_good_broadcast_unary_normal,