[Doc] Fix docstrings (#1216)

* fix docstrings

* remove fixme

* fix more docstring

* fix

Author: HydrogenSulfate

ppsci/arch/cuboid_transformer_decoder.py

@@ -756,7 +756,7 @@ class Upsample3DLayer(nn.Layer):
     Args:
         dim (int): The dimension of the input tensor.
         out_dim (int): The dimension of the output tensor.
-        target_size (Tuple[int,...]): The size of output tensor.
+        target_size (Tuple[int, ...]): The size of output tensor.
         temporal_upsample (bool, optional): Whether the temporal axis will go through upsampling. Defaults to False.
         kernel_size (int, optional): The kernel size of the Conv2D layer. Defaults to 3.
         layout (str, optional): The layout of the inputs. Defaults to "THWC".
@@ -856,7 +856,7 @@ class CuboidTransformerDecoder(nn.Layer):
 
     Args:
         target_temporal_length (int): The temporal length of the target.
-        mem_shapes (Tuple[int,...]): The mem shapes of the decoder.
+        mem_shapes (Tuple[int, ...]): The mem shapes of the decoder.
         cross_start (int, optional): The block to start cross attention. Defaults to 0.
         depth (list, optional): The number of layers for each block. Defaults to [2, 2].
         upsample_type (str, optional): The type of upsample. Defaults to "upsample".

ppsci/arch/cuboid_transformer_encoder.py

@@ -244,14 +244,14 @@ def update_cuboid_size_shift_size(data_shape, cuboid_size, shift_size, strategy)
     """Update the cuboid_size and shift_size
 
     Args:
-        data_shape (Tuple[int,...]): The shape of the data.
-        cuboid_size (Tuple[int,...]): Size of the cuboid.
-        shift_size (Tuple[int,...]): Size of the shift.
+        data_shape (Tuple[int, ...]): The shape of the data.
+        cuboid_size (Tuple[int, ...]): Size of the cuboid.
+        shift_size (Tuple[int, ...]): Size of the shift.
         strategy (str): The strategy of attention.
 
     Returns:
-        new_cuboid_size (Tuple[int,...]): Size of the cuboid.
-        new_shift_size (Tuple[int,...]): Size of the shift.
+        new_cuboid_size (Tuple[int, ...]): Size of the cuboid.
+        new_shift_size (Tuple[int, ...]): Size of the shift.
     """
 
     new_cuboid_size = list(cuboid_size)
@@ -271,8 +271,8 @@ def cuboid_reorder(data, cuboid_size, strategy):
 
     Args:
         data (paddle.Tensor): The input data.
-        cuboid_size (Tuple[int,...]): The size of the cuboid.
-        strategy (Tuple[int,...]): The cuboid strategy.
+        cuboid_size (Tuple[int, ...]): The size of the cuboid.
+        strategy (Tuple[int, ...]): The cuboid strategy.
 
     Returns:
         reordered_data (paddle.Tensor): Shape will be (B, num_cuboids, bT * bH * bW, C).
@@ -400,7 +400,9 @@ def masked_softmax(att_score, mask, axis: int = -1):
             att_score = att_score.masked_fill(paddle.logical_not(mask), -1e4)
         else:
             att_score = att_score.masked_fill(paddle.logical_not(mask), -1e18)
-        att_weights = nn.functional.softmax(x=att_score, axis=axis) * mask.astype(att_score.dtype)
+        att_weights = nn.functional.softmax(x=att_score, axis=axis) * mask.astype(
+            att_score.dtype
+        )
     else:
         att_weights = nn.functional.softmax(x=att_score, axis=axis)
     return att_weights
@@ -411,9 +413,9 @@ def cuboid_reorder_reverse(data, cuboid_size, strategy, orig_data_shape):
 
     Args:
         data (paddle.Tensor): The input data.
-        cuboid_size (Tuple[int,...]): The size of cuboid.
+        cuboid_size (Tuple[int, ...]): The size of cuboid.
         strategy (str): The strategy of reordering.
-        orig_data_shape (Tuple[int,...]): The original shape of the data.
+        orig_data_shape (Tuple[int, ...]): The original shape of the data.
 
     Returns:
         data (paddle.Tensor): The recovered data
@@ -1237,7 +1239,7 @@ class CuboidTransformerEncoder(nn.Layer):
     x --> attn_block --> patch_merge --> attn_block --> patch_merge --> ... --> out
 
     Args:
-        input_shape (Tuple[int,...]): The shape of the input. Contains T, H, W, C
+        input_shape (Tuple[int, ...]): The shape of the input. Contains T, H, W, C
         base_units (int, optional): The number of units. Defaults to 128.
         block_units (int, optional): The number of block units. Defaults to None.
         scale_alpha (float, optional):  We scale up the channels based on the formula:

ppsci/arch/cuboid_transformer_utils.py

@@ -82,7 +82,7 @@ def get_norm_layer(
 
     Args:
         normalization (str): The type of the layer normalization from ['layer_norm'].
-        axis (float): The axis to normalize the.
+        axis (int): The axis to normalize the.
         epsilon (float): The epsilon of the normalization layer.
         in_channels (int): Input channel.
 
@@ -236,12 +236,12 @@ def axial(self, input_shape):
         """Axial attention proposed in https://arxiv.org/abs/1912.12180
 
         Args:
-            input_shape (Tuple[int,...]): The shape of the input tensor, T H W.
+            input_shape (Tuple[int, ...]): The shape of the input tensor, T H W.
 
         Returns:
-            cuboid_size (Tuple[int,...]): The size of cuboid.
-            strategy (Tuple[str,...]): The strategy of the attention.
-            shift_size (Tuple[int,...]): The shift size of the attention.
+            cuboid_size (Tuple[int, ...]): The size of cuboid.
+            strategy (Tuple[str, ...]): The strategy of the attention.
+            shift_size (Tuple[int, ...]): The shift size of the attention.
         """
 
         T, H, W, _ = input_shape

ppsci/arch/cvit.py

@@ -504,11 +504,11 @@ def dot_product_attention_weights(
     you can directly call this function and call einsum yourself.
 
     Args:
-        query: queries for calculating attention with shape of [batch..., q_length,
+        query (paddle.Tensor): queries for calculating attention with shape of [batch..., q_length,
             num_heads, qk_depth_per_head].
-        key: keys for calculating attention with shape of [batch..., kv_length,
+        key (paddle.Tensor): keys for calculating attention with shape of [batch..., kv_length,
             num_heads, qk_depth_per_head].
-        bias: bias for the attention weights. This should be broadcastable to the
+        bias (Optional[paddle.Tensor]): bias for the attention weights. This should be broadcastable to the
             shape [batch..., num_heads, q_length, kv_length]. This can be used for
             incorporating causal masks, padding masks, proximity bias, etc.
 
@@ -555,13 +555,13 @@ def dot_product_attention(
     Note: query, key, value needn't have any batch dimensions.
 
     Args:
-        query: queries for calculating attention with shape of [batch..., q_length,
+        query (paddle.Tensor): queries for calculating attention with shape of [batch..., q_length,
             num_heads, qk_depth_per_head].
-        key: keys for calculating attention with shape of [batch..., kv_length,
+        key (paddle.Tensor): keys for calculating attention with shape of [batch..., kv_length,
             num_heads, qk_depth_per_head].
-        value: values to be used in attention with shape of [batch..., kv_length,
+        value (paddle.Tensor): values to be used in attention with shape of [batch..., kv_length,
             num_heads, v_depth_per_head].
-        bias: bias for the attention weights. This should be broadcastable to the
+        bias (Optional[paddle.Tensor]): bias for the attention weights. This should be broadcastable to the
             shape [batch..., num_heads, q_length, kv_length]. This can be used for
             incorporating causal masks, padding masks, proximity bias, etc.
 

ppsci/arch/dgmr.py

@@ -226,8 +226,8 @@ def forward(
         Perform the sampling from Skillful Nowcasting with GANs
 
         Args:
-            conditioning_states: Outputs from the `ContextConditioningStack` with the 4 input states, ordered from largest to smallest spatially
-            latent_dim: Output from `LatentConditioningStack` for input into the ConvGRUs
+            conditioning_states (List[paddle.Tensor]): Outputs from the `ContextConditioningStack` with the 4 input states, ordered from largest to smallest spatially
+            latent_dim (paddle.Tensor): Output from `LatentConditioningStack` for input into the ConvGRUs
         Returns:
             forecast_steps-length output of images for future timesteps
 
@@ -909,8 +909,10 @@ def __init__(
     def forward(self, x: paddle.Tensor) -> paddle.Tensor:
         """
         Args:
-            x: tensor on the correct device, to move over the latent distribution
-        Returns: z
+            x (paddle.Tensor): tensor on the correct device, to move over the latent distribution
+
+        Returns:
+            z
         """
         z = self.distribution.sample(self.shape)
         z = paddle.transpose(x=z, perm=(3, 0, 1, 2)).astype(dtype=x.dtype)

ppsci/arch/extformer_moe_cuboid.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Sequence
 from typing import Tuple
 from typing import Union
@@ -921,7 +923,7 @@ def get_initial_z(self, final_mem, T_out):
             raise NotImplementedError
         return initial_z
 
-    def forward(self, x: "paddle.Tensor", verbose: bool = False) -> "paddle.Tensor":
+    def forward(self, x: paddle.Tensor, verbose: bool = False) -> paddle.Tensor:
         """
         Args:
             x (paddle.Tensor): Tensor with shape (B, T, H, W, C).

ppsci/arch/extformer_moe_cuboid_decoder.py

@@ -855,7 +855,7 @@ class Upsample3DLayer(nn.Layer):
     Args:
         dim (int): The dimension of the input tensor.
         out_dim (int): The dimension of the output tensor.
-        target_size (Tuple[int,...]): The size of output tensor.
+        target_size (Tuple[int, ...]): The size of output tensor.
         temporal_upsample (bool, optional): Whether the temporal axis will go through upsampling. Defaults to False.
         kernel_size (int, optional): The kernel size of the Conv2D layer. Defaults to 3.
         layout (str, optional): The layout of the inputs. Defaults to "THWC".
@@ -956,7 +956,7 @@ class CuboidTransformerDecoder(nn.Layer):
 
     Args:
         target_temporal_length (int): The temporal length of the target.
-        mem_shapes (Tuple[int,...]): The mem shapes of the decoder.
+        mem_shapes (Tuple[int, ...]): The mem shapes of the decoder.
         cross_start (int, optional): The block to start cross attention. Defaults to 0.
         depth (list, optional): The number of layers for each block. Defaults to [2, 2].
         upsample_type (str, optional): The type of upsample. Defaults to "upsample".

ppsci/arch/extformer_moe_cuboid_encoder.py

@@ -268,14 +268,14 @@ def update_cuboid_size_shift_size(data_shape, cuboid_size, shift_size, strategy)
     """Update the cuboid_size and shift_size
 
     Args:
-        data_shape (Tuple[int,...]): The shape of the data.
-        cuboid_size (Tuple[int,...]): Size of the cuboid.
-        shift_size (Tuple[int,...]): Size of the shift.
+        data_shape (Tuple[int, ...]): The shape of the data.
+        cuboid_size (Tuple[int, ...]): Size of the cuboid.
+        shift_size (Tuple[int, ...]): Size of the shift.
         strategy (str): The strategy of attention.
 
     Returns:
-        new_cuboid_size (Tuple[int,...]): Size of the cuboid.
-        new_shift_size (Tuple[int,...]): Size of the shift.
+        new_cuboid_size (Tuple[int, ...]): Size of the cuboid.
+        new_shift_size (Tuple[int, ...]): Size of the shift.
     """
 
     new_cuboid_size = list(cuboid_size)
@@ -295,8 +295,8 @@ def cuboid_reorder(data, cuboid_size, strategy):
 
     Args:
         data (paddle.Tensor): The input data.
-        cuboid_size (Tuple[int,...]): The size of the cuboid.
-        strategy (Tuple[int,...]): The cuboid strategy.
+        cuboid_size (Tuple[int, ...]): The size of the cuboid.
+        strategy (Tuple[int, ...]): The cuboid strategy.
 
     Returns:
         reordered_data (paddle.Tensor): Shape will be (B, num_cuboids, bT * bH * bW, C).
@@ -437,9 +437,9 @@ def cuboid_reorder_reverse(data, cuboid_size, strategy, orig_data_shape):
 
     Args:
         data (paddle.Tensor): The input data.
-        cuboid_size (Tuple[int,...]): The size of cuboid.
+        cuboid_size (Tuple[int, ...]): The size of cuboid.
         strategy (str): The strategy of reordering.
-        orig_data_shape (Tuple[int,...]): The original shape of the data.
+        orig_data_shape (Tuple[int, ...]): The original shape of the data.
 
     Returns:
         data (paddle.Tensor): The recovered data
@@ -1391,7 +1391,7 @@ class CuboidTransformerEncoder(nn.Layer):
     x --> attn_block --> patch_merge --> attn_block --> patch_merge --> ... --> out
 
     Args:
-        input_shape (Tuple[int,...]): The shape of the input. Contains T, H, W, C
+        input_shape (Tuple[int, ...]): The shape of the input. Contains T, H, W, C
         base_units (int, optional): The number of units. Defaults to 128.
         block_units (int, optional): The number of block units. Defaults to None.
         scale_alpha (float, optional):  We scale up the channels based on the formula:

ppsci/arch/extformer_moe_cuboid_utils.py

@@ -82,7 +82,7 @@ def get_norm_layer(
 
     Args:
         normalization (str): The type of the layer normalization from ['layer_norm'].
-        axis (float): The axis to normalize the.
+        axis (int): The axis to normalize the.
         epsilon (float): The epsilon of the normalization layer.
         in_channels (int): Input channel.
 
@@ -239,12 +239,12 @@ def axial(self, input_shape):
         """Axial attention proposed in https://arxiv.org/abs/1912.12180
 
         Args:
-            input_shape (Tuple[int,...]): The shape of the input tensor, T H W.
+            input_shape (Tuple[int, ...]): The shape of the input tensor, T H W.
 
         Returns:
-            cuboid_size (Tuple[int,...]): The size of cuboid.
-            strategy (Tuple[str,...]): The strategy of the attention.
-            shift_size (Tuple[int,...]): The shift size of the attention.
+            cuboid_size (Tuple[int, ...]): The size of cuboid.
+            strategy (Tuple[str, ...]): The strategy of the attention.
+            shift_size (Tuple[int, ...]): The shift size of the attention.
         """
 
         T, H, W, _ = input_shape

ppsci/data/dataset/atmospheric_dataset.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from typing import List
 from typing import NamedTuple
 from typing import Optional
@@ -212,9 +214,9 @@ def stacked_to_dataset(
     All variables must have preserved_dims dimensions.
 
     Args:
-        stacked_array: Data in BHWC layout, encoded the same as dataset_to_stacked would if it was asked to encode `template_dataset`.
-        template_dataset: A template Dataset (or other mapping of DataArrays) demonstrating the shape of output required (variables, shapes, coordinates etc).
-        preserved_dims: dimensions from the target_template that were not folded in the predictions channels. The preserved_dims need to be a subset of the dims of all the variables of template_dataset.
+        stacked_array (xarray.Variable): Data in BHWC layout, encoded the same as dataset_to_stacked would if it was asked to encode `template_dataset`.
+        template_dataset (xarray.Dataset): A template Dataset (or other mapping of DataArrays) demonstrating the shape of output required (variables, shapes, coordinates etc).
+        preserved_dims (Tuple[str, ...]): dimensions from the target_template that were not folded in the predictions channels. The preserved_dims need to be a subset of the dims of all the variables of template_dataset.
 
     Returns:
         An xarray.Dataset (or other mapping of DataArrays) with the same shape and type as template_dataset.