Added ArcTan2 operation (#1079)

* Added ArcTan2 operation

* Cleanup, bug fixes from code review

* Minor cleanup, fixed Linux tests

mlx/backend/metal/kernels/binary.h

@@ -264,3 +264,10 @@ struct RightShift {
     return x >> y;
   };
 };
+
+struct ArcTan2 {
+  template <typename T>
+  T operator()(T y, T x) {
+    return metal::precise::atan2(y, x);
+  }
+};

mlx/backend/metal/kernels/binary.metal

@@ -241,6 +241,7 @@ instantiate_binary_types(mul, Multiply)
 instantiate_binary_types(sub, Subtract)
 instantiate_binary_types(pow, Power)
 instantiate_binary_types(rem, Remainder)
+instantiate_binary_float(arctan2, ArcTan2)
 
 // NaNEqual only needed for floating point types with boolean output
 instantiate_binary_all(naneq, float16, half, bool, NaNEqual)

mlx/backend/metal/primitives.cpp

@@ -451,6 +451,10 @@ void ArcTan::eval_gpu(const std::vector<array>& inputs, array& out) {
   unary_op(inputs, out, "arctan");
 }
 
+void ArcTan2::eval_gpu(const std::vector<array>& inputs, array& out) {
+  binary_op(inputs, out, "arctan2");
+}
+
 void ArcTanh::eval_gpu(const std::vector<array>& inputs, array& out) {
   unary_op(inputs, out, "arctanh");
 }