Сhange type of a tensor with bools (pytorch#19097)

Summary:
**This is **bc-breaking** change**
Change dtype of a tensor which was created from bool data.
Old behavior: torch.tensor([True, False]) -> uint8 tensor
Now: torch.tensor([True, False]) -> bool tensor

Tested via tests.
Pull Request resolved: pytorch#19097

Reviewed By: ezyang

Differential Revision: D15632553

Pulled By: izdeby

fbshipit-source-id: b019150844c561a6845710a3c62b12f06b68bbe3

Author: izdeby

aten/src/ATen/Declarations.cwrap

@@ -365,6 +365,7 @@
   name: _th_equal
   cpu_bool: True
   cname: equal
+  cpu_bool: True
   variants:
     - function
   return: bool

aten/src/ATen/core/jit_type.h

@@ -1254,7 +1254,7 @@ inline at::ScalarType scalarTypeFromJitType(const c10::TypePtr& type) {
   } else if (type == IntType::get()) {
     return at::ScalarType::Long;
   } else if (type == BoolType::get()) {
-    return at::ScalarType::Byte;
+    return at::ScalarType::Bool;
   }
   AT_ASSERTM(
       0,

test/test_dataloader.py

@@ -1007,7 +1007,7 @@ def test_default_collate_dtype(self):
         arr = [True, False]
         collated = _utils.collate.default_collate(arr)
         self.assertEqual(collated, torch.tensor(arr))
-        self.assertEqual(collated.dtype, torch.uint8)
+        self.assertEqual(collated.dtype, torch.bool)
 
         # Should be a no-op
         arr = ['a', 'b', 'c']

test/test_jit.py

@@ -6723,7 +6723,7 @@ def func7():
             return torch.tensor([[1]])
 
         list_input = [func1, func2, func3, func4, func5, func6, func7]
-        expected_shape = ["Long(*)", ("Byte(*)"), "Double(*)", "Double()", "Long()", "Byte()", "Long(*, *)"]
+        expected_shape = ["Long(*)", ("Bool(*)"), "Double(*)", "Double()", "Long()", "Bool()", "Long(*, *)"]
 
         for fn, expect in zip(list_input, expected_shape):
             self.checkScript(fn, ())

test/test_torch.py

@@ -3272,7 +3272,7 @@ def test_inference(default_dtype):
             self.assertIs(default_dtype, torch.tensor(()).dtype)
             self.assertIs(default_dtype, torch.tensor(5.).dtype)
             self.assertIs(torch.int64, torch.tensor(5).dtype)
-            self.assertIs(torch.uint8, torch.tensor(True).dtype)
+            self.assertIs(torch.bool, torch.tensor(True).dtype)
             self.assertIs(torch.int32, torch.tensor(5, dtype=torch.int32).dtype)
             self.assertIs(default_dtype, torch.tensor(((7, 5), (9, 5.))).dtype)
             self.assertIs(default_dtype, torch.tensor(((5., 5), (3, 5))).dtype)

torch/csrc/autograd/python_variable_indexing.cpp

@@ -173,7 +173,11 @@ static Variable applySlicing(const Variable& self, PyObject* index, variable_lis
           result = applySelect(result, dim, THPUtils_unpackLong(obj), i);
         } else {
           result = result.unsqueeze(dim);
-          handle_var(boolToIndexingTensor(result, var.item<uint8_t>() != 0));
+          if(scalar_type == at::kBool) {
+            handle_var(boolToIndexingTensor(result, var.item<bool>() != 0));
+          } else {
+            handle_var(boolToIndexingTensor(result, var.item<uint8_t>() != 0));
+          }
         }
       } else {
         handle_var(var);

torch/csrc/jit/register_special_ops.cpp

@@ -99,7 +99,6 @@ void storeLastDimension(
   }
 }
 
-// bool vector needs to be cast to uint8_t
 template <>
 void storeLastDimension<bool>(
     char* data,
@@ -112,7 +111,7 @@ void storeLastDimension<bool>(
   auto seq_size = obj.size();
   checkSequenceSize(n, dim, seq_size);
   for (int64_t i = 0; i < n; i++) {
-    *(uint8_t*)data = static_cast<uint8_t>(obj[i]);
+    *(bool*)data = static_cast<bool>(obj[i]);
     data += strides[dim] * elementSize;
   }
 }
@@ -291,7 +290,7 @@ RegisterOperators reg({
             DEFINE_TORCH_TENSOR_OP(
                 bool,
                 bool,
-                at::empty({}, at::CPU(at::kByte).options()).fill_(scalar_val))
+                at::empty({}, at::CPU(at::kBool).options()).fill_(scalar_val))
 
     // reference python implementation: internal_new_from_data in
     // tensor_new.cpp

torch/csrc/utils/tensor_new.cpp

@@ -124,8 +124,7 @@ ScalarType infer_scalar_type(PyObject *obj) {
     return ScalarType::Long;
   }
   if (PyBool_Check(obj)) {
-    // TODO: infer Bool when we have Bool ScalarType
-    return ScalarType::Byte;
+    return ScalarType::Bool;
   }
   if (THPVariable_Check(obj)) {
     auto var = reinterpret_cast<THPVariable*>(obj)->cdata;
@@ -477,7 +476,7 @@ Tensor indexing_tensor_from_data(
   // Specific to tensor indexing, converts an indexing list to an
   // indexing tensor (type Byte or Long)
   ScalarType inferred_scalar_type = infer_scalar_type(data);
-  if (inferred_scalar_type == ScalarType::Byte) {
+  if (inferred_scalar_type == ScalarType::Byte || inferred_scalar_type == ScalarType::Bool) {
     auto& idx_type = type.toScalarType(inferred_scalar_type);
     return internal_new_from_data(idx_type, inferred_scalar_type, std::move(device), data, false, false, false);
   } else {