pytorch · driazati · Aug 27, 2018 · Aug 27, 2018 · Aug 28, 2018 · Sep 2, 2018
diff --git a/aten/src/ATen/core/ivalue.cpp b/aten/src/ATen/core/ivalue.cpp
@@ -6,9 +6,11 @@
   _(Tensor) \
   _(Double) \
   _(Int) \
+  _(Bool) \
   _(Tuple) \
   _(IntList) \
   _(DoubleList) \
+  _(BoolList) \
   _(String) \
   _(TensorList) \
   _(Blob) \

diff --git a/aten/src/ATen/core/ivalue.h b/aten/src/ATen/core/ivalue.h
@@ -74,6 +74,7 @@ struct C10_EXPORT Tuple : public List<IValue> {
 using IntList = List<int64_t>;
 using TensorList = List<at::Tensor>;
 using DoubleList = List<double>;
+using BoolList = List<bool>;
 using GenericList = List<IValue>;
 
 // IValue is the generic tagged union used by the interpreter to hold
@@ -88,9 +89,11 @@ using GenericList = List<IValue>;
   _(Tensor) \
   _(Double) \
   _(Int) \
+  _(Bool) \
   _(Tuple) \
   _(IntList) \
   _(DoubleList) \
+  _(BoolList) \
   _(String) \
   _(TensorList) \
   _(Blob) \
@@ -224,8 +227,6 @@ struct CAFFE2_API IValue final {
   // allow you to pass literals (3, 4) without ambiguity
   IValue(int32_t i)
   : IValue(static_cast<int64_t>(i)) {}
-  IValue(bool b)
-  : IValue(static_cast<int64_t>(b)) {}
 
   bool isInt() const { return Tag::Int == tag; }
 
@@ -234,6 +235,17 @@ struct CAFFE2_API IValue final {
     return payload.as_int;
   }
 
+  // Bool
+  IValue(bool b)
+  : tag(Tag::Bool), is_intrusive_ptr(false) {
+    payload.as_bool = b;
+  }
+   bool isBool() const { return Tag::Bool == tag; }
+   bool toBool() const {
+    AT_ASSERT(isBool());
+    return payload.as_bool;
+  }
+
   // IntList
   IValue(c10::intrusive_ptr<IntList> v);
   IValue(std::vector<int64_t> v);
@@ -251,6 +263,7 @@ struct CAFFE2_API IValue final {
 
   const std::vector<int64_t>& toIntListRef() const;
   const std::vector<double>& toDoubleListRef() const;
+  const std::vector<bool>& toBoolListRef() const;
   const std::vector<at::Tensor>& toTensorListRef() const;
   const std::vector<IValue>& toGenericListRef() const;
 
@@ -280,6 +293,19 @@ struct CAFFE2_API IValue final {
     return toIntrusivePtr<DoubleList>();
   }
 
+  // BoolList
+  IValue(c10::intrusive_ptr<BoolList> v);
+  IValue(std::vector<bool> v);
+  bool isBoolList() const { return Tag::BoolList == tag; }
+  c10::intrusive_ptr<BoolList> toBoolList() && {
+    AT_ASSERT(isBoolList());
+    return moveToIntrusivePtr<BoolList>();
+  }
+  c10::intrusive_ptr<BoolList> toBoolList() const & {
+    AT_ASSERT(isBoolList());
+    return toIntrusivePtr<BoolList>();
+  }
+
   //TensorList
   IValue(c10::intrusive_ptr<TensorList> v);
   IValue(std::vector<at::Tensor> v);
@@ -323,15 +349,16 @@ struct CAFFE2_API IValue final {
     }
   }
   bool isScalar() {
-    return isDouble() || isInt();
+    return isDouble() || isInt() || isBool();
   }
   at::Scalar toScalar() const {
     if(isDouble())
       return toDouble();
     else if(isInt())
       return toInt();
-    else
-      throw std::runtime_error("IValue is not a Scalar");
+    else if (isBool())
+      return int(toBool());
+    throw std::runtime_error("IValue is not a Scalar");
   }
 
   // for debugging
@@ -396,6 +423,7 @@ struct CAFFE2_API IValue final {
   union {
     int64_t as_int;
     double as_double;
+    bool as_bool;
     c10::intrusive_ptr_target* as_intrusive_ptr;
     World as_world;
   } payload;
@@ -419,15 +447,16 @@ DEFINE_TO(at::Tensor, toTensor)
 DEFINE_TO(c10::intrusive_ptr<Tuple>, toTuple)
 DEFINE_TO(double, toDouble)
 DEFINE_TO(int64_t, toInt)
+DEFINE_TO(bool, toBool)
 DEFINE_TO(c10::intrusive_ptr<DoubleList>, toDoubleList)
 DEFINE_TO(c10::intrusive_ptr<IntList>, toIntList)
 DEFINE_TO(c10::intrusive_ptr<TensorList>, toTensorList)
 DEFINE_TO(c10::intrusive_ptr<GenericList>, toGenericList)
 DEFINE_TO(c10::intrusive_ptr<ConstantString>, toString)
 DEFINE_TO(at::Scalar, toScalar)
-DEFINE_TO(bool, toInt)
 DEFINE_TO(std::vector<int64_t>, toIntListRef)
 DEFINE_TO(std::vector<double>, toDoubleListRef)
+DEFINE_TO(std::vector<bool>, toBoolListRef)
 DEFINE_TO(std::vector<at::Tensor>, toTensorListRef)
 DEFINE_TO(std::vector<IValue>, toGenericListRef)
 DEFINE_TO(World, toWorld)
@@ -490,6 +519,13 @@ inline IValue::IValue(c10::intrusive_ptr<DoubleList> v)
 inline IValue::IValue(std::vector<double> v)
 : IValue(DoubleList::create(std::move(v))) {}
 
+inline IValue::IValue(c10::intrusive_ptr<BoolList> v)
+: tag(Tag::BoolList), is_intrusive_ptr(true) {
+  payload.as_intrusive_ptr = v.release();
+}
+inline IValue::IValue(std::vector<bool> v)
+: IValue(BoolList::create(std::move(v))) {}
+
 inline IValue::IValue(c10::intrusive_ptr<TensorList> v)
 : tag(Tag::TensorList), is_intrusive_ptr(true) {
   payload.as_intrusive_ptr = v.release();
@@ -517,6 +553,10 @@ inline const std::vector<at::Tensor>& IValue::toTensorListRef() const {
   return toTensorList()->elements();
 }
 
+inline const std::vector<bool>& IValue::toBoolListRef() const {
+  return toBoolList()->elements();
+}
+
 inline const std::vector<IValue>& IValue::toGenericListRef() const {
   return toGenericList()->elements();
 }

diff --git a/test/expect/TestBatched.test_for.expect b/test/expect/TestBatched.test_for.expect
@@ -5,7 +5,7 @@ graph(%x.1_data : Dynamic
       %y_mask : Dynamic
       %y_dims : Dynamic) {
   %6 : int = prim::Constant[value=10]()
-  %7 : int = prim::Constant[value=1]()
+  %7 : bool = prim::Constant[value=1]()
   %x : Dynamic, %9 : Dynamic, %10 : Dynamic = prim::Loop(%6, %7, %x.1_data, %x.1_mask, %x.1_dims)
     block0(%loop_num : int, %5_data : Dynamic, %5_mask : Dynamic, %5_dims : Dynamic) {
       %15 : int = prim::Constant[value=1]()
@@ -14,7 +14,7 @@ graph(%x.1_data : Dynamic
       %data.1 : Dynamic = aten::add(%5_data, %y_data, %alpha)
       %mask : Dynamic = aten::mul(%5_mask, %y_mask)
       %dims : Dynamic = aten::__or__(%5_dims, %y_dims)
-      %21 : int = prim::Constant[value=1]()
+      %21 : bool = prim::Constant[value=1]()
       %data : Dynamic = aten::where(%mask, %data.1, %5_data)
       -> (%21, %data, %mask, %dims)
     }

diff --git a/test/expect/TestBatched.test_if_else.expect b/test/expect/TestBatched.test_if_else.expect
@@ -7,7 +7,7 @@ graph(%a.1_data : Dynamic
   %6 : Dynamic = aten::gt(%a.1_data, %b_data)
   %7 : Dynamic = aten::mul(%a.1_mask, %b_mask)
   %8 : Dynamic = aten::__or__(%a.1_dims, %b_dims)
-  %9 : int = prim::TensorToNum(%6)
+  %9 : bool = prim::TensorToBool(%6)
   %10 : int = prim::Constant[value=1]()
   %11 : Long() = prim::NumToTensor(%10)
   %alpha.1 : float = prim::TensorToNum(%11)
@@ -24,17 +24,17 @@ graph(%a.1_data : Dynamic
   %23 : Dynamic = aten::type_as(%7, %6)
   %cond_mask.1 : Dynamic = aten::mul(%6, %23)
   %25 : int = aten::dim(%cond_mask.1)
-  %26 : int = aten::eq(%25, %22)
+  %26 : bool = aten::eq(%25, %22)
   %cond_data : Dynamic, %cond_mask : Dynamic, %data : Dynamic = prim::If(%26)
     block0() {
       %30 : int = aten::dim(%data.1)
       %31 : int = aten::sub(%30, %22)
-      %32 : int = prim::Constant[value=1]()
+      %32 : bool = prim::Constant[value=1]()
       %data.3 : Dynamic = prim::Loop(%31, %32, %cond_mask.1)
         block0(%_ : int, %35 : Dynamic) {
           %36 : int = aten::dim(%35)
           %data.2 : Dynamic = aten::unsqueeze(%35, %36)
-          %38 : int = prim::Constant[value=1]()
+          %38 : bool = prim::Constant[value=1]()
           -> (%38, %data.2)
         }
       %cond_data.1 : Dynamic = aten::expand_as(%data.3, %data.1)

diff --git a/test/expect/TestBatched.test_if_else_with_scalar.expect b/test/expect/TestBatched.test_if_else_with_scalar.expect
@@ -8,7 +8,7 @@ graph(%a.1_data : Dynamic
   %7 : Float() = prim::NumToTensor(%6)
   %other : float = prim::TensorToNum(%7)
   %9 : Dynamic = aten::gt(%a.1_data, %other)
-  %10 : int = prim::TensorToNum(%9)
+  %10 : bool = prim::TensorToBool(%9)
   %11 : int = prim::Constant[value=1]()
   %12 : Long() = prim::NumToTensor(%11)
   %alpha.1 : float = prim::TensorToNum(%12)
@@ -25,17 +25,17 @@ graph(%a.1_data : Dynamic
   %24 : Dynamic = aten::type_as(%a.1_mask, %9)
   %cond_mask.1 : Dynamic = aten::mul(%9, %24)
   %26 : int = aten::dim(%cond_mask.1)
-  %27 : int = aten::eq(%26, %23)
+  %27 : bool = aten::eq(%26, %23)
   %cond_data : Dynamic, %cond_mask : Dynamic, %data : Dynamic = prim::If(%27)
     block0() {
       %31 : int = aten::dim(%data.1)
       %32 : int = aten::sub(%31, %23)
-      %33 : int = prim::Constant[value=1]()
+      %33 : bool = prim::Constant[value=1]()
       %data.3 : Dynamic = prim::Loop(%32, %33, %cond_mask.1)
         block0(%_ : int, %36 : Dynamic) {
           %37 : int = aten::dim(%36)
           %data.2 : Dynamic = aten::unsqueeze(%36, %37)
-          %39 : int = prim::Constant[value=1]()
+          %39 : bool = prim::Constant[value=1]()
           -> (%39, %data.2)
         }
       %cond_data.1 : Dynamic = aten::expand_as(%data.3, %data.1)

diff --git a/test/expect/TestBatched.test_if_noelse.expect b/test/expect/TestBatched.test_if_noelse.expect
@@ -7,7 +7,7 @@ graph(%a.1_data : Dynamic
   %6 : Dynamic = aten::gt(%a.1_data, %b_data)
   %7 : Dynamic = aten::mul(%a.1_mask, %b_mask)
   %8 : Dynamic = aten::__or__(%a.1_dims, %b_dims)
-  %9 : int = prim::TensorToNum(%6)
+  %9 : bool = prim::TensorToBool(%6)
   %10 : int = prim::Constant[value=1]()
   %11 : Long() = prim::NumToTensor(%10)
   %alpha : float = prim::TensorToNum(%11)
@@ -18,17 +18,17 @@ graph(%a.1_data : Dynamic
   %17 : Dynamic = aten::type_as(%7, %6)
   %cond_mask.1 : Dynamic = aten::mul(%6, %17)
   %19 : int = aten::dim(%cond_mask.1)
-  %20 : int = aten::eq(%19, %16)
+  %20 : bool = aten::eq(%19, %16)
   %cond_data : Dynamic, %cond_mask : Dynamic, %data : Dynamic = prim::If(%20)
     block0() {
       %24 : int = aten::dim(%data.1)
       %25 : int = aten::sub(%24, %16)
-      %26 : int = prim::Constant[value=1]()
+      %26 : bool = prim::Constant[value=1]()
       %data.3 : Dynamic = prim::Loop(%25, %26, %cond_mask.1)
         block0(%_ : int, %29 : Dynamic) {
           %30 : int = aten::dim(%29)
           %data.2 : Dynamic = aten::unsqueeze(%29, %30)
-          %32 : int = prim::Constant[value=1]()
+          %32 : bool = prim::Constant[value=1]()
           -> (%32, %data.2)
         }
       %cond_data.1 : Dynamic = aten::expand_as(%data.3, %data.1)

diff --git a/test/expect/TestBatched.test_if_noelse_with_scalar.expect b/test/expect/TestBatched.test_if_noelse_with_scalar.expect
@@ -8,7 +8,7 @@ graph(%a.1_data : Dynamic
   %7 : Float() = prim::NumToTensor(%6)
   %other : float = prim::TensorToNum(%7)
   %9 : Dynamic = aten::gt(%a.1_data, %other)
-  %10 : int = prim::TensorToNum(%9)
+  %10 : bool = prim::TensorToBool(%9)
   %11 : int = prim::Constant[value=1]()
   %12 : Long() = prim::NumToTensor(%11)
   %alpha : float = prim::TensorToNum(%12)
@@ -19,17 +19,17 @@ graph(%a.1_data : Dynamic
   %18 : Dynamic = aten::type_as(%a.1_mask, %9)
   %cond_mask.1 : Dynamic = aten::mul(%9, %18)
   %20 : int = aten::dim(%cond_mask.1)
-  %21 : int = aten::eq(%20, %17)
+  %21 : bool = aten::eq(%20, %17)
   %cond_data : Dynamic, %cond_mask : Dynamic, %data : Dynamic = prim::If(%21)
     block0() {
       %25 : int = aten::dim(%data.1)
       %26 : int = aten::sub(%25, %17)
-      %27 : int = prim::Constant[value=1]()
+      %27 : bool = prim::Constant[value=1]()
       %data.3 : Dynamic = prim::Loop(%26, %27, %cond_mask.1)
         block0(%_ : int, %30 : Dynamic) {
           %31 : int = aten::dim(%30)
           %data.2 : Dynamic = aten::unsqueeze(%30, %31)
-          %33 : int = prim::Constant[value=1]()
+          %33 : bool = prim::Constant[value=1]()
           -> (%33, %data.2)
         }
       %cond_data.1 : Dynamic = aten::expand_as(%data.3, %data.1)

diff --git a/test/expect/TestBatched.test_while.expect b/test/expect/TestBatched.test_while.expect
@@ -8,12 +8,12 @@ graph(%a.1_data : Dynamic
   %7 : Dynamic = aten::gt(%a.1_data, %b_data)
   %8 : Dynamic = aten::mul(%a.1_mask, %b_mask)
   %9 : Dynamic = aten::__or__(%a.1_dims, %b_dims)
-  %10 : int = prim::TensorToNum(%7)
+  %10 : bool = prim::TensorToBool(%7)
   %11 : int = prim::Constant[value=0]()
   %12 : Dynamic = aten::mul(%7, %8)
   %13 : Dynamic = aten::sum(%12)
   %14 : Dynamic = aten::gt(%13, %11)
-  %15 : int = prim::TensorToNum(%14)
+  %15 : bool = prim::TensorToBool(%14)
   %16 : Dynamic, %17 : Dynamic, %18 : Dynamic, %a : Dynamic, %20 : Dynamic, %21 : Dynamic = prim::Loop(%6, %15, %7, %8, %9, %a.1_data, %a.1_mask, %a.1_dims)
     block0(%loop_num : int, %cond_data.2 : Dynamic, %cond_mask.3 : Dynamic, %cond_dims : Dynamic, %6_data : Dynamic, %6_mask : Dynamic, %6_dims : Dynamic) {
       %29 : int = prim::Constant[value=1]()
@@ -25,22 +25,22 @@ graph(%a.1_data : Dynamic
       %35 : Dynamic = aten::gt(%data.1, %b_data)
       %36 : Dynamic = aten::mul(%mask, %b_mask)
       %37 : Dynamic = aten::__or__(%dims, %b_dims)
-      %38 : int = prim::TensorToNum(%35)
+      %38 : bool = prim::TensorToBool(%35)
       %39 : int = prim::Constant[value=1]()
       %40 : Dynamic = aten::type_as(%cond_mask.3, %cond_data.2)
       %cond_mask.1 : Dynamic = aten::mul(%cond_data.2, %40)
       %42 : int = aten::dim(%cond_mask.1)
-      %43 : int = aten::eq(%42, %39)
+      %43 : bool = aten::eq(%42, %39)
       %cond_data : Dynamic, %cond_mask : Dynamic, %data : Dynamic = prim::If(%43)
         block0() {
           %47 : int = aten::dim(%data.1)
           %48 : int = aten::sub(%47, %39)
-          %49 : int = prim::Constant[value=1]()
+          %49 : bool = prim::Constant[value=1]()
           %data.3 : Dynamic = prim::Loop(%48, %49, %cond_mask.1)
             block0(%_ : int, %52 : Dynamic) {
               %53 : int = aten::dim(%52)
               %data.2 : Dynamic = aten::unsqueeze(%52, %53)
-              %55 : int = prim::Constant[value=1]()
+              %55 : bool = prim::Constant[value=1]()
               -> (%55, %data.2)
             }
           %cond_data.1 : Dynamic = aten::expand_as(%data.3, %data.1)
@@ -57,7 +57,7 @@ graph(%a.1_data : Dynamic
       %62 : Dynamic = aten::mul(%35, %36)
       %63 : Dynamic = aten::sum(%62)
       %64 : Dynamic = aten::gt(%63, %61)
-      %65 : int = prim::TensorToNum(%64)
+      %65 : bool = prim::TensorToBool(%64)
       -> (%65, %35, %36, %37, %res_data, %res_mask, %res_dims)
     }
   return (%a, %20, %21);

diff --git a/test/expect/TestJit.test_batchnorm.expect b/test/expect/TestJit.test_batchnorm.expect
@@ -4,10 +4,10 @@ graph(%0 : Double(2, 2, 2, 2)
       %3 : Double(2)
       %4 : Double(2)
       %5 : Long()) {
-  %6 : int = prim::Constant[value=1](), scope: BatchNorm2d
+  %6 : bool = prim::Constant[value=1](), scope: BatchNorm2d
   %7 : float = prim::Constant[value=0.1](), scope: BatchNorm2d
   %8 : float = prim::Constant[value=1e-05](), scope: BatchNorm2d
-  %9 : int = prim::Constant[value=1](), scope: BatchNorm2d
+  %9 : bool = prim::Constant[value=1](), scope: BatchNorm2d
   %10 : Double(2, 2, 2, 2) = aten::batch_norm(%0, %1, %2, %3, %4, %6, %7, %8, %9), scope: BatchNorm2d
   return (%10);
 }
diff --git a/test/expect/TestJit.test_constant_prop_if_constant.expect b/test/expect/TestJit.test_constant_prop_if_constant.expect
@@ -1,10 +1,10 @@
 graph(%a : Dynamic
       %b : Dynamic) {
   %c2.1 : int = prim::Constant[value=1]()
-  %3 : int = prim::TensorToNum(%a)
+  %3 : bool = prim::TensorToBool(%a)
   %c0.4 : int, %c1 : int = prim::If(%3)
     block0() {
-      %6 : int = prim::TensorToNum(%b)
+      %6 : bool = prim::TensorToBool(%b)
       %c0.3 : int = prim::If(%6)
         block0() {
           %8 : int = prim::Constant[value=2]()

diff --git a/test/expect/TestJit.test_constant_prop_loop_constant.expect b/test/expect/TestJit.test_constant_prop_loop_constant.expect
@@ -3,16 +3,16 @@ graph() {
   %b.2 : int = prim::Constant[value=1]()
   %2 : int = prim::Constant[value=2147483647]()
   %b.1 : int = prim::Constant[value=0]()
-  %4 : int = prim::Constant[value=1]()
+  %4 : bool = prim::Constant[value=1]()
   %b.3 : int = prim::Loop(%2, %4, %b.1)
     block0(%6 : int, %7 : int) {
-      %8 : int = prim::Constant[value=1]()
+      %8 : bool = prim::Constant[value=1]()
       -> (%8, %b.2)
     }
-  %9 : int = prim::Constant[value=0]()
+  %9 : bool = prim::Constant[value=0]()
   %b : int = prim::Loop(%2, %9, %b.3)
     block0(%11 : int, %12 : int) {
-      %13 : int = prim::Constant[value=0]()
+      %13 : bool = prim::Constant[value=0]()
       -> (%13, %b.4)
     }
   return (%b);