where() infer type of the other condition when one of the two is a te… (

jonasrauber#39)

* where() infer type of the other condition when one of the two is a tensor.

* correct flake8

eagerpy/tensor/jax.py

@@ -489,6 +489,9 @@ def tanh(self: TensorType) -> TensorType:
     def float32(self: TensorType) -> TensorType:
         return self.astype(np.float32)
 
+    def float64(self: TensorType) -> TensorType:
+        return self.astype(np.float32)
+
     def where(self: TensorType, x: TensorOrScalar, y: TensorOrScalar) -> TensorType:
         x, y = unwrap_(x, y)
         return type(self)(np.where(self.raw, x, y))

eagerpy/tensor/numpy.py

@@ -394,6 +394,9 @@ def tanh(self: TensorType) -> TensorType:
     def float32(self: TensorType) -> TensorType:
         return self.astype(np.float32)
 
+    def float64(self: TensorType) -> TensorType:
+        return self.astype(np.float64)
+
     def where(self: TensorType, x: TensorOrScalar, y: TensorOrScalar) -> TensorType:
         x, y = unwrap_(x, y)
         return type(self)(np.where(self.raw, x, y))

eagerpy/tensor/pytorch.py

@@ -535,19 +535,31 @@ def sqrt(self: TensorType) -> TensorType:
     def float32(self: TensorType) -> TensorType:
         return self.astype(torch.float32)
 
+    def float64(self: TensorType) -> TensorType:
+        return self.astype(torch.float64)
+
     def where(self: TensorType, x: TensorOrScalar, y: TensorOrScalar) -> TensorType:
+
         if isinstance(x, Tensor):
             x_ = x.raw
         elif isinstance(x, int) or isinstance(x, float):
-            x_ = torch.full_like(self.raw, x, dtype=torch.float32)
+            if isinstance(y, Tensor):
+                dtype = y.raw.dtype
+            else:
+                dtype = torch.float32
+            x_ = torch.full_like(self.raw, x, dtype=dtype)
         else:
             raise TypeError(
                 "expected x to be a Tensor, int or float"
             )  # pragma: no cover
         if isinstance(y, Tensor):
             y_ = y.raw
         elif isinstance(y, int) or isinstance(y, float):
-            y_ = torch.full_like(self.raw, y, dtype=torch.float32)
+            if isinstance(x, Tensor):
+                dtype = x.raw.dtype
+            else:
+                dtype = torch.float32
+            y_ = torch.full_like(self.raw, y, dtype=dtype)
         return type(self)(torch.where(self.raw, x_, y_))
 
     def __lt__(self: TensorType, other: TensorOrScalar) -> TensorType:

eagerpy/tensor/tensor.py

@@ -209,6 +209,10 @@ def tanh(self: TensorType) -> TensorType:
     def float32(self: TensorType) -> TensorType:
         ...
 
+    @abstractmethod
+    def float64(self: TensorType) -> TensorType:
+        ...
+
     @abstractmethod
     def where(self: TensorType, x: TensorOrScalar, y: TensorOrScalar) -> TensorType:
         ...

eagerpy/tensor/tensorflow.py

@@ -501,6 +501,9 @@ def tanh(self: TensorType) -> TensorType:
     def float32(self: TensorType) -> TensorType:
         return self.astype(tf.float32)
 
+    def float64(self: TensorType) -> TensorType:
+        return self.astype(tf.float64)
+
     def where(self: TensorType, x: TensorOrScalar, y: TensorOrScalar) -> TensorType:
         x, y = unwrap_(x, y)
         return type(self)(tf.where(self.raw, x, y))

tests/test_main.py

@@ -1248,6 +1248,18 @@ def test_where_second_scalar(t: Tensor) -> Tensor:
     return ep.where(t >= 3, t, 2)
 
 
+@compare_all
+def test_where_first_scalar64(dummy: Tensor) -> Tensor:
+    t = ep.arange(dummy, 60).float64().reshape((3, 4, 5))
+    return ep.where(t >= 3, 2, -t)
+
+
+@compare_all
+def test_where_second_scalar64(dummy: Tensor) -> Tensor:
+    t = ep.arange(dummy, 60).float64().reshape((3, 4, 5))
+    return ep.where(t >= 3, t, 2)
+
+
 @compare_all
 def test_where_both_scalars(t: Tensor) -> Tensor:
     return ep.where(t >= 3, 2, 5)