Merge pull request #3967 from kmaehashi/add-floordiv

implement floordiv and rfloordiv

chainer/functions/math/basic_math.py

@@ -4,6 +4,7 @@
 from chainer.backends import cuda
 from chainer import function_node
 import chainer.functions
+from chainer.functions.math import floor as _floor
 from chainer.functions.math import matmul as _matmul
 from chainer import utils
 from chainer.utils import type_check
@@ -445,6 +446,26 @@ def rdiv(self, rhs):  # rhs / lhs
     return DivFromConstant(rhs).apply((self,))[0]
 
 
+def floordiv(self, rhs):  # lhs // rhs
+    """Element-wise floor division.
+
+    Returns:
+        ~chainer.Variable: Output variable.
+    """
+
+    return _floor.floor(div(self, rhs))
+
+
+def rfloordiv(self, rhs):  # rhs // lhs
+    """Element-wise floor division.
+
+    Returns:
+        ~chainer.Variable: Output variable.
+    """
+
+    return _floor.floor(rdiv(self, rhs))
+
+
 class PowVarVar(function_node.FunctionNode):
 
     @property
@@ -846,6 +867,8 @@ def install_variable_arithmetics():
     variable.Variable.__truediv__ = div
     variable.Variable.__rdiv__ = rdiv
     variable.Variable.__rtruediv__ = rdiv
+    variable.Variable.__floordiv__ = floordiv
+    variable.Variable.__rfloordiv__ = rfloordiv
     variable.Variable.__pow__ = pow
     variable.Variable.__rpow__ = rpow
     variable.Variable.__matmul__ = matmul

tests/chainer_tests/functions_tests/math_tests/test_basic_math.py

@@ -50,6 +50,9 @@ def test_mul_forward_cpu(self):
     def test_div_forward_cpu(self):
         self.forward_cpu(lambda x, y: x / y)
 
+    def test_floordiv_forward_cpu(self):
+        self.forward_cpu(lambda x, y: x // y)
+
     def test_pow_forward_cpu(self):
         self.forward_cpu(lambda x, y: x ** y)
 
@@ -65,6 +68,9 @@ def test_rmul_forward_cpu(self):
     def test_rdiv_forward_cpu(self):
         self.forward_cpu(lambda x, y: y.__rtruediv__(x))
 
+    def test_rfloordiv_forward_cpu(self):
+        self.forward_cpu(lambda x, y: y.__rfloordiv__(x))
+
     def test_rpow_forward_cpu(self):
         self.forward_cpu(lambda x, y: y.__rpow__(x))
 
@@ -87,6 +93,10 @@ def test_mul_forward_gpu(self):
     def test_div_forward_gpu(self):
         self.forward_gpu(lambda x, y: x / y)
 
+    @attr.gpu
+    def test_floordiv_forward_gpu(self):
+        self.forward_gpu(lambda x, y: x // y)
+
     @attr.gpu
     def test_pow_forward_gpu(self):
         self.forward_gpu(lambda x, y: x ** y)
@@ -107,6 +117,10 @@ def test_rmul_forward_gpu(self):
     def test_rdiv_forward_gpu(self):
         self.forward_gpu(lambda x, y: y.__rtruediv__(x))
 
+    @attr.gpu
+    def test_rfloordiv_forward_gpu(self):
+        self.forward_gpu(lambda x, y: y.__rfloordiv__(x))
+
     @attr.gpu
     def test_rpow_forward_gpu(self):
         self.forward_gpu(lambda x, y: y.__rpow__(x))