Add support for __floordiv__ and __rdiv__ for integral tensors

apaszke · apaszke · commit da65cb2e329b · 2018-05-03T18:39:11.000+02:00
diff --git a/torch/tensor.py b/torch/tensor.py
@@ -317,7 +317,11 @@ def __rsub__(self, other):
         return -self + other
 
     def __rdiv__(self, other):
-        return self.reciprocal() * other
+        if self.dtype.is_floating_point:
+            return self.reciprocal() * other
+        else:
+            return (self.double().reciprocal() * other).type_as(self)
+
     __rtruediv__ = __rdiv__
     __itruediv__ = _C._TensorBase.__idiv__
 
@@ -334,6 +338,18 @@ def __ipow__(self, other):
     def __rpow__(self, other):
         return self.new([other]) ** self
 
+    def __floordiv__(self, other):
+        result = self / other
+        if result.dtype.is_floating_point:
+            result = result.trunc()
+        return result
+
+    def __rfloordiv__(self, other):
+        result = other / self
+        if result.dtype.is_floating_point:
+            result = result.trunc()
+        return result
+
     __neg__ = _C._TensorBase.neg
 
     __eq__ = _C._TensorBase.eq
