start adding rules to propagate scalar to results

Author: zdevito

Formatting.cpp

@@ -230,9 +230,9 @@ std::ostream& print(std::ostream& stream, const Tensor & tensor_, int64_t linesi
     stream << "[ Tensor (empty) ]";
   } else {
     Tensor tensor = tensor_.toType(getType(kCPU,kDouble)).contiguous();
-    if(tensor.ndimension() == 0) {
+    if(tensor_.ndimension() == 0) {
       stream << std::defaultfloat << tensor.data<double>()[0] << std::endl;
-      stream << "[" << tensor_.pImpl->toString() << " of size {}]";
+      stream << "[ " << tensor_.pImpl->toString() << "{} ]";
     } else if(tensor.ndimension() == 1) {
       double scale;
       int64_t sz;
@@ -244,17 +244,17 @@ std::ostream& print(std::ostream& stream, const Tensor & tensor_, int64_t linesi
       for(int64_t i = 0; i < tensor.size(0); i++) {
         stream << std::setw(sz) << tensor_p[i]/scale << std::endl;
       }
-      stream << "[" << tensor_.pImpl->toString() << " of size " << tensor.size(0) << "]";
+      stream << "[ " << tensor_.pImpl->toString() << "{" << tensor.size(0) << "} ]";
     } else if(tensor.ndimension() == 2) {
       __printMatrix(stream, tensor, linesize, 0);
-      stream << "[" << tensor_.pImpl->toString() <<" of size " << tensor.size(0) << "x" <<  tensor.size(1) << "]";
+      stream << "[ " << tensor_.pImpl->toString() << "{" << tensor.size(0) << "," <<  tensor.size(1) << "} ]";
     } else {
         __printTensor(stream, tensor, linesize);
-        stream << "[" << tensor_.pImpl->toString() << " of size " << tensor.size(0);
+        stream << "[ " << tensor_.pImpl->toString() << "{" << tensor.size(0);
         for(int64_t i = 1; i < tensor.ndimension(); i++) {
-          stream << "x" << tensor.size(i);
+          stream << "," << tensor.size(i);
         }
-        stream << "]";
+        stream << "} ]";
     }
   }
   return stream;

TensorImpl.h

@@ -38,8 +38,14 @@ struct TensorImpl {
   bool isScalar() const {
     return is_scalar;
   }
-  void setScalar(bool s) {
-    is_scalar = s;
+  // this is called by the generated wrapper code when there are conditions
+  // when this output tensor should be a scalar. e.g. when all inputs
+  // to a function 'add' were scalars, then condition_when_scalar == true.
+  // we also prevent this from getting marked as a scalar if it is not
+  // the right shape afterall.
+  TensorImpl* maybeScalar(bool condition_when_scalar) {
+    is_scalar = condition_when_scalar && (dim() == 0 || dim() == 1 && sizes()[0] == 1);
+    return this;
   }
 
 private:

copy_wrapper.py

@@ -32,7 +32,7 @@
       throw std::runtime_error("unsupported type in copy");
       break;
   }
-  dst.pImpl->setScalar(src.pImpl->isScalar());
+  dst.pImpl->maybeScalar(src.pImpl->isScalar());
 }
 """)
 

function_wrapper.py

@@ -318,10 +318,12 @@ def emit_body(env, option):
         # referencing another
         seen_names = set()
         count = 0
+        scalar_check = None
         for arg in option['arguments']:
             if is_real_argument_to_wrapper(arg):
                 count += 1
-
+            if arg['type'] == 'THSize*':
+                scalar_check = '{}.size() == 0'.format(arg['name'])
             # only generated checked casts the first time we see it
             if not arg['name'] in seen_names and requires_checked_cast(arg):
                 seen_names.add(arg['name'])
@@ -369,24 +371,32 @@ def emit_body(env, option):
 
         call = prefix + CodeTemplate("${cname}(${derived_actuals})").substitute(env)
         ret = option['return']
+
         if ret['kind'] == 'arguments':
             body.append(call + ";")
             arguments_indices = ret['arguments']
+            arguments = [option['arguments'][argi]
+                         for argi in arguments_indices]
+            if scalar_check is not None:
+                for arg in arguments:
+                    body.append("bool maybe_scalar = {};".format(scalar_check))
+                    body.append("{}_->maybeScalar(maybe_scalar);".format(arg['name']))
             if len(arguments_indices) == 1:
-                arg = option['arguments'][arguments_indices[0]]
+                arg = arguments[0]
                 body.append("return {};".format(arg['name']))
             else:
-                arguments = [option['arguments'][argi]
-                             for argi in arguments_indices]
                 types = [to_return_type(arg, option) for arg in arguments]
                 # TODO: check for move semantics...
                 names = [arg['name'] for arg in arguments]
                 body.append(CodeTemplate("return std::tuple<${types}>(${names});").substitute(
                     types=types, names=names))
         elif ret['kind'] == 'type':
             if ret['type'] == 'THTensor*':
+                maybe_scalar = "->maybeScalar({})".format(scalar_check) \
+                               if scalar_check is not None \
+                               else ""
                 body.append(CodeTemplate(
-                    "return Tensor(new ${Tensor}(context,${arg_name}),false);").substitute(env, arg_name=call))
+                    "return Tensor((new ${Tensor}(context,${arg_name}))${maybe_scalar},false);").substitute(env, arg_name=call,maybe_scalar=maybe_scalar))
             else:
                 # we using int64_t for long in the API, so correct it here...
                 if is_actual_return_long(ret):

test/basic.cpp

@@ -172,9 +172,6 @@ static void test(Type & type) {
       std::cout << "zero-dim: " << std::endl;
       Tensor a =  type.scalarTensor(4); //type.rand({1});
 
-      // TODO: automate
-      a.pImpl->setScalar(true);
-
       std::cout << a << "dims: " << a.dim() << std::endl;
       std::cout << Scalar(a) << std::endl;
       Tensor b = type.rand({3,4});