!27581 Add bool for tensorarray

Merge pull request !27581 from VectorSL/add-bool-for-tensorarray

mindspore/nn/reinforcement/tensor_array.py

@@ -28,13 +28,13 @@ class TensorArray(Cell):
 
     Args:
         dtype (mindspore.dtype): the data type in the TensorArray.
-        element_shape (List[int]): the shape of each tensor in a TensorArray.
+        element_shape (tuple[int]): the shape of each tensor in a TensorArray.
         dynamic_size (bool): if true, the size of TensorArray can be increased. Default: True.
         size (int): if dynamic_size=False, `size` means the max_size of the TensorArray.
         name (string): the name of this TensorArray. Default: "TA".
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -58,7 +58,7 @@ class TensorArray(Cell):
     def __init__(self, dtype, element_shape, dynamic_size=True, size=0, name="TA"):
         """Initialize TensorArray"""
         super(TensorArray, self).__init__()
-        Validator.check_subclass("dtype", dtype, mstype.number_type, self.cls_name)
+        Validator.check_subclass("dtype", dtype, mstype.number_type + (mstype.bool_,), self.cls_name)
         Validator.check_int(size, 0, Rel.GE, "size", self.cls_name)
         self.handle_ = ta.TensorArray(dtype, element_shape, dynamic_size, size, name)()
         self.tensor_array_write = ta.TensorArrayWrite()

mindspore/ops/operations/_tensor_array.py

@@ -27,7 +27,7 @@ class TensorArray(PrimitiveWithInfer):
 
     Args:
         dtype (mindspore.dtype): the data type in the TensorArray.
-        element_shape (List[int]): the shape of each tensor in a TensorArray.
+        element_shape (tuple[int]): the shape of each tensor in a TensorArray.
         dynamic_size (bool): If true the TensorArray can increase the size. Default: True.
         size (int): The size of the TensorArray if dynamic_size = False.
         name (string): the name of this TensorArray. Default: "TA".
@@ -39,7 +39,7 @@ class TensorArray(PrimitiveWithInfer):
         - **output** (Tensor[mindspore.int64]) - an unique handle binded to the TensorArray.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -51,7 +51,7 @@ class TensorArray(PrimitiveWithInfer):
     """
     @prim_attr_register
     def __init__(self, dtype, element_shape, dynamic_size=True, size=0, name="TA"):
-        validator.check_type_name("dtype", dtype, mstype.number_type, self.name)
+        validator.check_type_name("dtype", dtype, mstype.number_type + (mstype.bool_,), self.name)
         validator.check_int(size, 0, Rel.GE, "size", self.name)
         self.add_prim_attr('dtype', dtype)
         self.add_prim_attr('element_shape', element_shape)
@@ -79,7 +79,7 @@ class TensorArrayWrite(PrimitiveWithInfer):
         None.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -99,7 +99,7 @@ class TensorArrayWrite(PrimitiveWithInfer):
     def infer_dtype(self, handle_type, index_type, value_type):
         validator.check_type_name("handle", handle_type, (ms.int64), self.name)
         validator.check_type_name("index", index_type, (int, ms.int64), self.name)
-        validator.check_type_name("value", value_type, mstype.number_type, self.name)
+        validator.check_type_name("value", value_type, mstype.number_type + (mstype.bool_,), self.name)
         return mstype.int64
 
 class TensorArrayRead(PrimitiveWithInfer):
@@ -108,7 +108,7 @@ class TensorArrayRead(PrimitiveWithInfer):
 
     Args:
         dtype (mindspore.dtype): the data type in the TensorArray.
-        element_shape (List[int]): the shape of each tensor in a TensorArray.
+        element_shape (tuple[int]): the shape of each tensor in a TensorArray.
 
     Inputs:
         - **index** (Tensor[int64]) - The position to read.
@@ -118,7 +118,7 @@ class TensorArrayRead(PrimitiveWithInfer):
         - **output** (Tensor) - the value in position index.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -134,7 +134,7 @@ class TensorArrayRead(PrimitiveWithInfer):
     """
     @prim_attr_register
     def __init__(self, dtype, element_shape):
-        validator.check_type_name("dtype", dtype, mstype.number_type, self.name)
+        validator.check_type_name("dtype", dtype, mstype.number_type + (mstype.bool_,), self.name)
         self.add_prim_attr('dtype', dtype)
         self.add_prim_attr('element_shape', element_shape)
         self.add_prim_attr('side_effect_mem', True)
@@ -160,7 +160,7 @@ class TensorArrayClose(PrimitiveWithInfer):
         None.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -192,7 +192,7 @@ class TensorArrayClear(PrimitiveWithInfer):
         None.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -219,7 +219,7 @@ class TensorArrayStack(Primitive):
 
     Args:
         dtype (mindspore.dtype): the data type in the TensorArray.
-        element_shape (List[int]): the shape of each tensor in a TensorArray.
+        element_shape (tuple[int]): the shape of each tensor in a TensorArray.
 
     Inputs:
         - **handle** (mindspore.int64) - The handle pointed to the TensorArray.
@@ -228,7 +228,7 @@ class TensorArrayStack(Primitive):
         - **output** (Tensor) - the stacked value from the TensorArray.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -263,7 +263,7 @@ class TensorArraySize(PrimitiveWithInfer):
         - **output** (Tensor[mindspore.int64]) - the logical size of the TensorArray.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore

tests/st/ops/cpu/test_tensor_array.py

@@ -84,3 +84,11 @@ def test_tensorarray():
     expect_s = [0., 0., 1., 0., 0., 1.]
     assert np.allclose(s.asnumpy(), expect_s)
     ta.close()
+    ta = nn.TensorArray(mindspore.bool_, ())
+    ta.write(1, Tensor(True, mindspore.bool_))
+    s = ta.stack()
+    v = ta.read(1)
+    expect_s = [False, True]
+    assert np.allclose(v.asnumpy(), expect_s[1])
+    assert np.allclose(s.asnumpy(), expect_s)
+    ta.close()

tests/st/ops/gpu/test_tensor_array.py

@@ -91,3 +91,11 @@ def test_tensorarray():
     expect_s = [0., 0., 1., 0., 0., 1.]
     assert np.allclose(s.asnumpy(), expect_s)
     ta.close()
+    ta = nn.TensorArray(mindspore.bool_, ())
+    ta.write(1, Tensor(True, mindspore.bool_))
+    s = ta.stack()
+    v = ta.read(1)
+    expect_s = [False, True]
+    assert np.allclose(v.asnumpy(), expect_s[1])
+    assert np.allclose(s.asnumpy(), expect_s)
+    ta.close()