167 tensor redistribution for binary operations

heat/core/arithmetics.py

@@ -1,8 +1,10 @@
 import torch
+import numpy as np
 
 from .communication import MPI
 from . import operations
 from . import dndarray
+from . import factories
 
 
 __all__ = [
@@ -128,7 +130,25 @@ def fmod(t1, t2):
     tensor([[0., 0.]
             [2., 2.]])
     """
-    return operations.__binary_op(torch.fmod, t1, t2)
+
+    if np.isscalar(t1):
+    #Special treatment for fmod, since torch operation fmod only supports formats (Tensor input, Tensor other, Tensor out) or (Tensor input, Number other, Tensor out)
+        if np.isscalar(t2):
+            try:
+                tensor_1 = factories.array(t1)
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+        else:
+            try:
+                tensor_1 = factories.array([t1])
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+
+        result = operations.__binary_op(torch.fmod, tensor_1, t2)
+        return result
+    else:
+        return operations.__binary_op(torch.fmod, t1, t2)
+
 
 
 def mod(t1, t2):
@@ -247,7 +267,15 @@ def pow(t1, t2):
     tensor([[1., 8.],
             [27., 64.]])
     """
-    return operations.__binary_op(torch.pow, t1, t2)
+
+    if np.isscalar(t1) and np.isscalar(t2):
+        try:
+            tensor_1 = factories.array(t1)
+        except (ValueError, TypeError,):
+            raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+        return operations.__binary_op(torch.pow, tensor_1, t2)
+    else:
+        return operations.__binary_op(torch.pow, t1, t2)
 
 
 def sub(t1, t2):

heat/core/operations.py

@@ -36,49 +36,60 @@ def __binary_op(operation, t1, t2):
     """
 
     if np.isscalar(t1):
-        try:
-            t1 = factories.array([t1])
-        except (ValueError, TypeError,):
-            raise TypeError('Data type not supported, input was {}'.format(type(t1)))
-
         if np.isscalar(t2):
             try:
-                t2 = factories.array([t2])
+                result = operation(t1, t2)
             except (ValueError, TypeError,):
-                raise TypeError('Only numeric scalars are supported, but input was {}'.format(type(t2)))
+                raise TypeError('Only numeric scalars are supported, but inputs were {} and  {}'.format(type(t1), type(t2)))
             output_shape = (1,)
+            output_type = types.canonical_heat_type(result.dtype)
             output_split = None
             output_device = None
             output_comm = MPI_WORLD
+
         elif isinstance(t2, dndarray.DNDarray):
+            try:
+                result = operation(t1, t2._DNDarray__array)
+            except (ValueError, TypeError,):
+                raise TypeError('Data type not supported, input was {}'.format(type(t1)))
+
             output_shape = t2.shape
+            output_type = types.canonical_heat_type(result.dtype)
             output_split = t2.split
             output_device = t2.device
             output_comm = t2.comm
+
         else:
             raise TypeError('Only tensors and numeric scalars are supported, but input was {}'.format(type(t2)))
 
-        if t1.dtype != t2.dtype:
-            t1 = t1.astype(t2.dtype)
 
     elif isinstance(t1, dndarray.DNDarray):
         if np.isscalar(t2):
             try:
-                t2 = factories.array([t2])
-                output_shape = t1.shape
-                output_split = t1.split
-                output_device = t1.device
-                output_comm = t1.comm
+                result = operation(t1._DNDarray__array, t2)
             except (ValueError, TypeError,):
                 raise TypeError('Data type not supported, input was {}'.format(type(t2)))
 
+            output_shape = t1.shape
+            output_type = types.canonical_heat_type(result.dtype)
+            output_split = t1.split
+            output_device = t1.device
+            output_comm = t1.comm
+
         elif isinstance(t2, dndarray.DNDarray):
+
+            if t2.dtype != t1.dtype:
+                t2 = t2.astype(t1.dtype)
+
+            output_shape = stride_tricks.broadcast_shape(t1.shape, t2.shape)
+
             # TODO: implement complex NUMPY rules
-            if t2.split is None or t2.split == t1.split:
-                output_shape = stride_tricks.broadcast_shape(t1.shape, t2.shape)
-                output_split = t1.split
-                output_device = t1.device
-                output_comm = t1.comm
+            if t2.split == t1.split:
+                pass
+            elif (t1.split is not None) and (t2.split is None):
+                t2.resplit(axis=t1.split)
+            elif (t2.split is not None) and (t1.split is None):
+                t1.resplit(axis=t2.split)
             else:
                 # It is NOT possible to perform binary operations on tensors with different splits, e.g. split=0
                 # and split=1
@@ -99,30 +110,33 @@ def __binary_op(operation, t1, t2):
                         t2._DNDarray__array = torch.zeros(t2.shape, dtype=t2.dtype.torch_type())
                     t2.comm.Bcast(t2)
 
+            promoted_type = types.promote_types(t1.dtype, t2.dtype).torch_type()
+            if t1.split is not None:
+                if len(t1.lshape) > t1.split and t1.lshape[t1.split] == 0:
+                    result = t1._DNDarray__array.type(promoted_type)
+                else:
+                    result = operation(t1._DNDarray__array.type(promoted_type), t2._DNDarray__array.type(promoted_type))
+            elif t2.split is not None:
+                if len(t2.lshape) > t2.split and t2.lshape[t2.split] == 0:
+                    result = t2._DNDarray__array.type(promoted_type)
+                else:
+                    result = operation(t1._DNDarray__array.type(promoted_type), t2._DNDarray__array.type(promoted_type))
+            else:
+                result = operation(t1._DNDarray__array.type(promoted_type), t2._DNDarray__array.type(promoted_type))
+
+            output_type = types.canonical_heat_type(t1.dtype)
+            output_split = t1.split
+            output_device = t1.device
+            output_comm = t1.comm
+
         else:
             raise TypeError('Only tensors and numeric scalars are supported, but input was {}'.format(type(t2)))
 
-        if t2.dtype != t1.dtype:
-            t2 = t2.astype(t1.dtype)
-
     else:
-        raise NotImplementedError('Not implemented for non scalar')
+        raise TypeError('Only tensors and numeric scalars are supported, but input was {}'.format(type(t1)))
 
-    promoted_type = types.promote_types(t1.dtype, t2.dtype).torch_type()
-    if t1.split is not None:
-        if len(t1.lshape) > t1.split and t1.lshape[t1.split] == 0:
-            result = t1._DNDarray__array.type(promoted_type)
-        else:
-            result = operation(t1._DNDarray__array.type(promoted_type), t2._DNDarray__array.type(promoted_type))
-    elif t1.split is not None:
-        if len(t2.lshape) > t2.split and t2.lshape[t2.split] == 0:
-            result = t2._DNDarray__array.type(promoted_type)
-        else:
-            result = operation(t1._DNDarray__array.type(promoted_type), t2._DNDarray__array.type(promoted_type))
-    else:
-        result = operation(t1._DNDarray__array.type(promoted_type), t2._DNDarray__array.type(promoted_type))
 
-    return dndarray.DNDarray(result, output_shape, types.canonical_heat_type(t1.dtype), output_split, output_device, output_comm)
+    return dndarray.DNDarray(result, output_shape, output_type, output_split, output_device, output_comm)
 
 
 def __local_op(operation, x, out, **kwargs):

heat/core/relational.py

@@ -48,7 +48,22 @@ def eq(t1, t2):
     tensor([[0, 1],
             [0, 0]])
     """
-    return operations.__binary_op(torch.eq, t1, t2)
+    if np.isscalar(t1):
+        if np.isscalar(t2):
+            try:
+                tensor_1 = factories.array(t1)
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+        else:
+            try:
+                tensor_1 = factories.array([t1])
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+
+        result = operations.__binary_op(torch.eq, tensor_1, t2)
+        return result
+    else:
+        return operations.__binary_op(torch.eq, t1, t2)
 
 
 def equal(t1, t2):
@@ -80,7 +95,24 @@ def equal(t1, t2):
     >>> ht.eq(T1, 3.0)
     False
     """
-    result_tensor = operations.__binary_op(torch.equal, t1, t2)
+    if np.isscalar(t1):
+        try:
+            tensor_1 = factories.array([t1])
+        except (ValueError, TypeError,):
+            raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+    else:
+        tensor_1 = t1
+    if np.isscalar(t2):
+        try:
+            tensor_2 = factories.array([t2])
+        except (ValueError, TypeError,):
+            raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+    else:
+        tensor_2 = t2
+
+    result_tensor = operations.__binary_op(torch.equal, tensor_1, tensor_2)
+
+
     result_value = result_tensor._DNDarray__array
     if isinstance(result_value, torch.Tensor):
         result_value = True
@@ -120,7 +152,22 @@ def ge(t1, t2):
     tensor([[0, 1],
             [1, 1]], dtype=torch.uint8)
     """
-    return operations.__binary_op(torch.ge, t1, t2)
+    if np.isscalar(t1):
+        if np.isscalar(t2):
+            try:
+                tensor_1 = factories.array(t1)
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+        else:
+            try:
+                tensor_1 = factories.array([t1])
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+
+        result = operations.__binary_op(torch.ge, tensor_1, t2)
+        return result
+    else:
+        return operations.__binary_op(torch.ge, t1, t2)
 
 
 def gt(t1, t2):
@@ -156,7 +203,22 @@ def gt(t1, t2):
     tensor([[0, 0],
             [1, 1]], dtype=torch.uint8)
     """
-    return operations.__binary_op(torch.gt, t1, t2)
+    if np.isscalar(t1):
+        if np.isscalar(t2):
+            try:
+                tensor_1 = factories.array(t1)
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+        else:
+            try:
+                tensor_1 = factories.array([t1])
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+
+        result = operations.__binary_op(torch.gt, tensor_1, t2)
+        return result
+    else:
+        return operations.__binary_op(torch.gt, t1, t2)
 
 
 def le(t1, t2):
@@ -191,7 +253,22 @@ def le(t1, t2):
     tensor([[1, 1],
             [0, 0]], dtype=torch.uint8)
     """
-    return operations.__binary_op(torch.le, t1, t2)
+    if np.isscalar(t1):
+        if np.isscalar(t2):
+            try:
+                tensor_1 = factories.array(t1)
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+        else:
+            try:
+                tensor_1 = factories.array([t1])
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+
+        result = operations.__binary_op(torch.le, tensor_1, t2)
+        return result
+    else:
+        return operations.__binary_op(torch.le, t1, t2)
 
 
 def lt(t1, t2):
@@ -226,7 +303,22 @@ def lt(t1, t2):
     tensor([[1, 0],
             [0, 0]], dtype=torch.uint8)
     """
-    return operations.__binary_op(torch.lt, t1, t2)
+    if np.isscalar(t1):
+        if np.isscalar(t2):
+            try:
+                tensor_1 = factories.array(t1)
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+        else:
+            try:
+                tensor_1 = factories.array([t1])
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+
+        result = operations.__binary_op(torch.lt, tensor_1, t2)
+        return result
+    else:
+        return operations.__binary_op(torch.lt, t1, t2)
 
 
 def ne(t1, t2):
@@ -260,4 +352,19 @@ def ne(t1, t2):
     tensor([[1, 0],
             [1, 1]])
     """
-    return operations.__binary_op(torch.ne, t1, t2)
+    if np.isscalar(t1):
+        if np.isscalar(t2):
+            try:
+                tensor_1 = factories.array(t1)
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+        else:
+            try:
+                tensor_1 = factories.array([t1])
+            except (ValueError, TypeError,):
+                raise TypeError('First operand must be numeric scalar or NDNArray, but was {}'.format(type(t1)))
+
+        result = operations.__binary_op(torch.ne, tensor_1, t2)
+        return result
+    else:
+        return operations.__binary_op(torch.ne, t1, t2)