Add the remainder op (#85)

* Add remainder in the C++ backend
* Add the python binding and test

mlx/backend/accelerate/primitives.cpp

@@ -322,6 +322,45 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
   }
 }
 
+// TODO: Avoid code duplication with the common backend.
+struct RemainderFn {
+  template <typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return std::fmod(numerator, denominator);
+  }
+
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return numerator % denominator;
+  }
+};
+
+void Remainder::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+
+  if (a.dtype() == float32) {
+    binary(
+        a,
+        b,
+        out,
+        RemainderFn{},
+        UseDefaultBinaryOp(),
+        UseDefaultBinaryOp(),
+        [](const auto* a, const auto* b, auto* o, auto n) {
+          int num_el = n;
+          vvremainderf((float*)o, (const float*)a, (const float*)b, &num_el);
+        });
+  } else {
+    binary(a, b, out, RemainderFn{});
+  }
+}
+
 void Exp::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];

mlx/backend/common/binary.cpp

@@ -82,6 +82,29 @@ void Divide::eval(const std::vector<array>& inputs, array& out) {
   binary(a, b, out, [](auto x, auto y) { return x / y; });
 }
 
+struct RemainderFn {
+  template <typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return std::fmod(numerator, denominator);
+  }
+
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return numerator % denominator;
+  }
+};
+
+void Remainder::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary(a, b, out, RemainderFn{});
+}
+
 void Equal::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   if (equal_nan_) {

mlx/backend/common/default_primitives.cpp

@@ -35,6 +35,7 @@ DEFAULT(Copy)
 DEFAULT(Cos)
 DEFAULT(Cosh)
 DEFAULT(Divide)
+DEFAULT(Remainder)
 DEFAULT(Equal)
 DEFAULT(Erf)
 DEFAULT(ErfInv)

mlx/backend/metal/kernels/binary.metal

@@ -14,6 +14,13 @@ struct Divide {
   template <typename T> T operator()(T x, T y) { return x / y; }
 };
 
+struct Remainder {
+  template <typename T> T operator()(T x, T y) { return x % y; }
+  template <> float operator()(float x, float y) { return fmod(x, y); }
+  template <> half operator()(half x, half y) { return fmod(x, y); }
+  template <> bfloat operator()(bfloat x, bfloat y) { return fmod(x, y); }
+};
+
 struct Equal {
   template <typename T> bool operator()(T x, T y) { return x == y; }
 };
@@ -363,6 +370,7 @@ instantiate_binary_types(min, Minimum)
 instantiate_binary_types(mul, Multiply)
 instantiate_binary_types(sub, Subtract)
 instantiate_binary_types(pow, Power)
+instantiate_binary_types(rem, Remainder)
 
 // NaNEqual only needed for floating point types with boolean output
 instantiate_binary_all(naneq, float16, half, bool, NaNEqual)

mlx/backend/metal/kernels/complex.h

@@ -110,3 +110,7 @@ constexpr complex64_t operator-(complex64_t a, complex64_t b) {
 constexpr complex64_t operator*(complex64_t a, complex64_t b) {
   return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real};
 }
+
+constexpr complex64_t operator%(complex64_t a, complex64_t b) {
+  return {fmod(a.real, b.real), fmod(a.imag, b.imag)};
+}

mlx/backend/metal/primitives.cpp

@@ -363,6 +363,10 @@ void Divide::eval_gpu(const std::vector<array>& inputs, array& out) {
   binary_op(inputs, out, "div");
 }
 
+void Remainder::eval_gpu(const std::vector<array>& inputs, array& out) {
+  binary_op(inputs, out, "rem");
+}
+
 void Equal::eval_gpu(const std::vector<array>& inputs, array& out) {
   binary_op(inputs, out, equal_nan_ ? "naneq" : "eq");
 }

mlx/backend/no_metal/primitives.cpp

@@ -30,6 +30,7 @@ NO_GPU(Copy)
 NO_GPU(Cos)
 NO_GPU(Cosh)
 NO_GPU(Divide)
+NO_GPU(Remainder)
 NO_GPU(Equal)
 NO_GPU(Erf)
 NO_GPU(ErfInv)