diff --git a/mindspore/nn/reinforcement/tensor_array.py b/mindspore/nn/reinforcement/tensor_array.py index 9367639c2e6..0b837a28d5e 100644 --- a/mindspore/nn/reinforcement/tensor_array.py +++ b/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() diff --git a/mindspore/ops/operations/_tensor_array.py b/mindspore/ops/operations/_tensor_array.py index d2788d5413d..8cc7aa60224 100644 --- a/mindspore/ops/operations/_tensor_array.py +++ b/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 diff --git a/tests/st/ops/cpu/test_tensor_array.py b/tests/st/ops/cpu/test_tensor_array.py index 3c92431322e..175baeeebc4 100644 --- a/tests/st/ops/cpu/test_tensor_array.py +++ b/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() diff --git a/tests/st/ops/gpu/test_tensor_array.py b/tests/st/ops/gpu/test_tensor_array.py index a99cf24e521..1b23c64030e 100644 --- a/tests/st/ops/gpu/test_tensor_array.py +++ b/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()