Add broadcasting support for map_, map2_.

gchanan · soumith · commit 5c56c72a143e · 2017-06-07T19:07:04.000-04:00
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -985,20 +985,22 @@ def _test_broadcast(self, cast):
             "dist", "atan2", "pow", "lerp", "add",
             "sub", "mul", "div", "fmod", "remainder",
             "eq", "ge", "gt", "le", "lt", "max", "min", "ne",
-            "addcdiv", "addcmul", "masked_copy", "masked_fill"
+            "addcdiv", "addcmul", "masked_copy", "masked_fill",
+            "map", "map2"
         ]
         # functions with no torch. equivalent
-        fns_no_torch = ["sub", "masked_copy", "masked_fill"]
+        fns_no_torch = ["sub", "masked_copy", "masked_fill", "map", "map2"]
         # functions with no inplace equivalent
         fns_no_inplace = ["dist", "max", "min"]
         # functions with no out-of-place tensor version
-        fns_no_out_place = ["masked_copy", "masked_fill"]
+        fns_no_out_place = ["masked_copy", "masked_fill", "map", "map2"]
         # functions with fallback to equal nElem behavior
         fns_fallback = ["add", "sub", "div", "mul", "pow", "fmod", "remainder",
                         "eq", "ge", "gt", "le", "lt", "max", "min", "ne",
-                        "addcdiv", "addcmul", "masked_copy", "masked_fill"]
+                        "addcdiv", "addcmul", "masked_copy", "masked_fill",
+                        "map", "map2"]
         # functions with three tensor arguments
-        fns_3_args = ["addcdiv", "addcmul"]
+        fns_3_args = ["addcdiv", "addcmul", "map2"]
 
         for fn in fns:
             (dims_small, dims_large, dims_full) = TestTorch._select_broadcastable_dims(self)
@@ -1091,6 +1093,10 @@ def tensorfn_inplace(t0, t1, t2=None):
                         return t0_fn(t1 < 0.5, t1.expand_as(t0))
                     elif fn == "masked_fill":
                         return t0_fn(t1 < 0.5, 1.0)
+                    elif fn == "map":
+                        return t0_fn(t1, lambda x, y: x + y)
+                    elif fn == "map2":
+                        return t0_fn(t1, t2, lambda x, y,z : x + y + z)
                     elif fn in fns_3_args:
                         return t0_fn(1.0, t1, t2)
                     else:
@@ -1134,13 +1140,14 @@ def _test_broadcast_fallback(self, cast):
         # functions that should fallback to pointwise behavior
         fns_fallback = ["add", "sub", "div", "mul", "pow", "fmod", "remainder",
                         "eq", "ge", "gt", "le", "lt", "max", "min", "ne",
-                        "addcdiv", "addcmul", "masked_copy", "masked_fill"]
+                        "addcdiv", "addcmul", "masked_copy", "masked_fill",
+                        "map", "map2"]
         # functions with three tensor arguments
-        fns_3_args = ["addcdiv", "addcmul"]
+        fns_3_args = ["addcdiv", "addcmul", "map2"]
         # functions with no inplace equivalent
         fns_no_inplace = ["max", "min"]
         # functions with no out-of-place tensor version
-        fns_no_out_place = ["masked_copy", "masked_fill"]
+        fns_no_out_place = ["masked_copy", "masked_fill", "map", "map2"]
 
         for fn in fns_fallback:
             # case 1: both broadcastable and nElems equal -- verify that we broadcast
@@ -1162,6 +1169,10 @@ def tensorfn(myfn, t1, t2):
                         return myfn(t1 < 0.5, torch.randn(4*4).float())
                     elif fn == "masked_fill":
                         return myfn(t1 < 0.5, 1.0)
+                    elif fn == "map":
+                        return myfn(t1, lambda x, y: x + y)
+                    elif fn == "map2":
+                         return myfn(t1, t2, lambda x, y, z: x + y + z)
                     elif fn in fns_3_args:
                         return myfn(1.0, t1, t2)
                     else:
diff --git a/tools/cwrap/plugins/Broadcast.py b/tools/cwrap/plugins/Broadcast.py
@@ -55,7 +55,8 @@ def getPreArgStringTemplate(self, includeElementCount=True, type=None):
         if type == None:
             ret = """THTensor *${arg_op_other}_save = ${arg_op_other};
                      THTensorPtr ${arg_op_other}_guard = THTensor_(new)(LIBRARY_STATE_NOARGS);
-                     ${arg_op_other}=${arg_op_other}_guard.get();"""
+                     ${arg_op_other}=${arg_op_other}_guard.get();
+                  """
             if includeElementCount:
                 ret += "ptrdiff_t ${arg_op_other}_nElem = THTensor_(nElement)(LIBRARY_STATE ${arg_op_other}_save);"
         else:
@@ -165,9 +166,9 @@ def getInPlacePreExpand1Template(self, type=None):
         ret = """bool ${arg_op_other}_raise = ${raise_errors} || (${arg_op_a}_nElem != ${arg_op_other}_nElem);
                  int ${arg_op_other}_err ="""
         if type == None:
-            ret += """!skip_expand && THTensor_(expand)(LIBRARY_STATE"""
+            ret += """!skip_expand && THTensor_(expand)(LIBRARY_STATE\n"""
         else:
-            ret += """!skip_expand && THBroadcastTensor_(expand)(LIBRARY_STATE"""
+            ret += """!skip_expand && THBroadcastTensor_(expand)(LIBRARY_STATE\n"""
 
         ret += """                                               ${arg_op_other},
                                                                  ${arg_op_other}_save,
@@ -187,9 +188,9 @@ def getInPlacePreExpand2Template(self, type1=None, type2=None):
                  bool ${arg_op_other2}_raise = ${raise_errors} || (${arg_op_a}_nElem != ${arg_op_other2}_nElem);
                  int ${arg_op_other1}_err ="""
         if type1 is None:
-            ret += """!skip_expand && THTensor_(expand)(LIBRARY_STATE"""
+            ret += """!skip_expand && THTensor_(expand)(LIBRARY_STATE\n"""
         else:
-            ret += """!skip_expand && THBroadcastTensor_(expand)(LIBRARY_STATE"""
+            ret += """!skip_expand && THBroadcastTensor_(expand)(LIBRARY_STATE\n"""
 
         ret += """                                              ${arg_op_other1},
                                                                 ${arg_op_other1}_save,
@@ -202,9 +203,9 @@ def getInPlacePreExpand2Template(self, type1=None, type2=None):
         ret += """}
                int ${arg_op_other2}_err ="""
         if type2 == None:
-            ret += """!skip_expand && THTensor_(expand)(LIBRARY_STATE"""
+            ret += """!skip_expand && THTensor_(expand)(LIBRARY_STATE\n"""
         else:
-            ret += """!skip_expand && THBroadcastTensor_(expand)(LIBRARY_STATE"""
+            ret += """!skip_expand && THBroadcastTensor_(expand)(LIBRARY_STATE\n"""
         ret += """                                                ${arg_op_other2},
                                                                   ${arg_op_other2}_save,
                                                                   ${arg_op_a}_size.get(),
@@ -317,7 +318,7 @@ def process_option_code_template(self, template, option):
                         arg_op_other2=arg_op_c,
                         post_code=post_code)
                     expand_code += self.IN_PLACE_BACK_COMPAT_WARN_TEMPLATE.substitute(op_b_mapping)
-                    expand_code += self.IN_PLACE_BACK_COMPAT_WARN_TEMPLATE.substitute(op_b_mapping)
+                    expand_code += self.IN_PLACE_BACK_COMPAT_WARN_TEMPLATE.substitute(op_c_mapping)
                 else:
                     expand_code = self.getInPlacePreExpand1Template(type=type_op_b).substitute(op_b_mapping, post_code=post_code)
                     expand_code += self.IN_PLACE_BACK_COMPAT_WARN_TEMPLATE.substitute(op_b_mapping)
diff --git a/torch/_tensor_docs.py b/torch/_tensor_docs.py
@@ -792,8 +792,10 @@
 map_(tensor, callable)
 
 Applies :attr:`callable` for each element in this tensor and the given tensor
-and stores the results in this tensor. The :attr:`callable` should have the
-signature::
+and stores the results in this tensor.  This tensor and the given tensor must be
+:ref:`broadcastable <broadcasting-semantics>`.
+
+The :attr:`callable` should have the signature::
 
     def callable(a, b) -> number
 """)
diff --git a/torch/csrc/generic/methods/TensorApply.cwrap b/torch/csrc/generic/methods/TensorApply.cwrap
@@ -57,6 +57,36 @@ static PyObject * THPTensor_(map)(THPTensor *self, PyObject *args)
 
   THTensor *tensor = self->cdata;
   THTensor *src = src_object->cdata;
+
+  THLongStoragePtr tensor_size = THTensor_(newSizeOf)(LIBRARY_STATE tensor);
+  ptrdiff_t tensor_nElem = THTensor_(nElement)(LIBRARY_STATE tensor);
+  bool skip_expand = false;
+  THTensor *src_save = src;
+  THTensorPtr src_guard = THTensor_(new)(LIBRARY_STATE_NOARGS);
+  src=src_guard.get();
+  ptrdiff_t src_nElem = THTensor_(nElement)(LIBRARY_STATE src_save);
+
+  bool src_raise = false || (tensor_nElem != src_nElem);
+  int src_err =!skip_expand && THTensor_(expand)(LIBRARY_STATE
+        src,
+        src_save,
+        tensor_size.get(),
+        src_raise);
+  if (src_err != 0 && !src_raise) {
+    skip_expand = true; // don't do further expansions
+    src = src_save;
+    PyErr_WarnEx(PyExc_UserWarning, "src is not broadcastable to tensor, but they have the same number of "
+        "elements.  Falling back to deprecated pointwise behavior.", 1);
+  }
+  if (getBackCompatBroadcastWarn()) {
+    bool same_shape = THSize_isSameSizeAs(tensor->size, tensor->nDimension,
+        src_save->size, src_save->nDimension);
+    if (!same_shape && src_err == 0 && (tensor_nElem == src_nElem) && !false) {
+      PyErr_WarnEx(PyExc_UserWarning, "tensor and src do not have the same shape, but are broadcastable, and have the same number of "
+          "elements.  Changing behavior in a backwards incompatible manner to broadcasting rather than viewing as 1-dimensional.", 1);
+    }
+  }
+
   TH_TENSOR_APPLY2(real, tensor, real, src,
                   PyObject *ret =
                       PyObject_CallFunction(fn, (char*)(BUILD_REAL_FMT BUILD_REAL_FMT),
@@ -71,6 +101,8 @@ static PyObject * THPTensor_(map)(THPTensor *self, PyObject *args)
                   Py_DECREF(ret);
                   );
 
+  src = src_save;
+
   Py_INCREF(self);
   return (PyObject*)self;
   END_HANDLE_TH_ERRORS
@@ -95,6 +127,62 @@ static PyObject * THPTensor_(map2)(THPTensor *self, PyObject *args)
   THTensor *tensor = self->cdata;
   THTensor *src1 = src1_object->cdata;
   THTensor *src2 = src2_object->cdata;
+
+
+  THLongStoragePtr tensor_size = THTensor_(newSizeOf)(LIBRARY_STATE tensor);
+  ptrdiff_t tensor_nElem = THTensor_(nElement)(LIBRARY_STATE tensor);
+  bool skip_expand = false;
+  THTensor *src1_save = src1;
+  THTensorPtr src1_guard = THTensor_(new)(LIBRARY_STATE_NOARGS);
+  src1=src1_guard.get();
+  ptrdiff_t src1_nElem = THTensor_(nElement)(LIBRARY_STATE src1_save);
+  THTensor *src2_save = src2;
+  THTensorPtr src2_guard = THTensor_(new)(LIBRARY_STATE_NOARGS);
+  src2=src2_guard.get();
+  ptrdiff_t src2_nElem = THTensor_(nElement)(LIBRARY_STATE src2_save);
+  bool src1_raise = false || (tensor_nElem != src1_nElem);
+  bool src2_raise = false || (tensor_nElem != src2_nElem);
+  int src1_err =!skip_expand && THTensor_(expand)(LIBRARY_STATE
+      src1,
+      src1_save,
+      tensor_size.get(),
+      src1_raise || src2_raise);
+  if (src1_err != 0 && !(src1_raise || src2_raise)) {
+    skip_expand = true; // don't do further expansions
+    src1 = src1_save;
+    src2 = src2_save;
+    PyErr_WarnEx(PyExc_UserWarning, "src1, src2 are not broadcastable to self, but they all have the same number of "
+        "elements.  Falling back to deprecated pointwise behavior.", 1);
+  }
+  int src2_err =!skip_expand && THTensor_(expand)(LIBRARY_STATE
+      src2,
+      src2_save,
+      tensor_size.get(),
+      src1_raise || src2_raise);
+  if (src2_err != 0 && !(src1_raise || src2_raise)) {
+    skip_expand = true; // don't do further expansions
+    src1 = src1_save;
+    src2 = src2_save;
+    PyErr_WarnEx(PyExc_UserWarning, "src1, src2 are not broadcastable to self, but they all have the same number of "
+        "elements.  Falling back to deprecated pointwise behavior.", 1);
+  }
+  if (getBackCompatBroadcastWarn()) {
+    bool same_shape = THSize_isSameSizeAs(tensor->size, tensor->nDimension,
+        src1_save->size, src1_save->nDimension);
+    if (!same_shape && src1_err == 0 && (tensor_nElem == src1_nElem) && !false) {
+      PyErr_WarnEx(PyExc_UserWarning, "tensor and src1 do not have the same shape, but are broadcastable, and have the same number of "
+          "elements.  Changing behavior in a backwards incompatible manner to broadcasting rather than viewing as 1-dimensional.", 1);
+    }
+  }
+  if (getBackCompatBroadcastWarn()) {
+    bool same_shape = THSize_isSameSizeAs(tensor->size, tensor->nDimension,
+        src1_save->size, src1_save->nDimension);
+    if (!same_shape && src1_err == 0 && (tensor_nElem == src1_nElem) && !false) {
+      PyErr_WarnEx(PyExc_UserWarning, "tensor and src2 do not have the same shape, but are broadcastable, and have the same number of "
+          "elements.  Changing behavior in a backwards incompatible manner to broadcasting rather than viewing as 1-dimensional.", 1);
+    }
+  }
+
   TH_TENSOR_APPLY3(real, tensor, real, src1, real, src2,
                   PyObject *ret =
                       PyObject_CallFunction(fn, (char*)(BUILD_REAL_FMT BUILD_REAL_FMT BUILD_REAL_FMT),
@@ -109,6 +197,9 @@ static PyObject * THPTensor_(map2)(THPTensor *self, PyObject *args)
                   Py_DECREF(ret);
                   );
 
+  src1 = src1_save;
+  src2 = src2_save;
+
   Py_INCREF(self);
   return (PyObject*)self;
   END_HANDLE_TH_ERRORS
