@@ -118,11 +118,11 @@ def _make_metadata(self, tensor: Optional[torch.Tensor]) -> Metadata:
118118 """
119119 Create the metadata as a dictionary based on the input tensor.
120120
121- Parameters :
122- - tensor: The input tensor or None if no tensor is provided.
121+ Args :
122+ tensor: The input tensor or None if no tensor is provided.
123123
124124 Returns:
125- - metadata: A dictionary with the following keys:
125+ metadata: A dictionary with the following keys:
126126 - "dtype": The data type of the tensor or None.
127127 - "shape": The shape of the tensor or None.
128128 """
@@ -135,13 +135,13 @@ def _prepare_recv_buffer(self, metadata: Metadata) -> torch.Tensor:
135135 """
136136 Create a buffer to receive the tensor based on the provided metadata.
137137
138- Parameters :
139- - metadata: A dictionary with keys "dtype" and "shape", describing
140- the tensor's data type and shape.
138+ Args :
139+ metadata: A dictionary with keys "dtype" and "shape",
140+ describing the tensor's data type and shape.
141141
142142 Returns:
143- - buffer: A tensor of the specified type and shape, allocated on
144- self.device.
143+ buffer: A tensor of the specified type and shape,
144+ allocated on ` self.device` .
145145 """
146146 return torch .empty (metadata ["shape" ],
147147 dtype = metadata ["dtype" ],
@@ -151,8 +151,8 @@ def _send_metadata(self, metadata: Metadata):
151151 """
152152 Send the metadata dictionary to the target rank.
153153
154- Parameters :
155- - metadata: A dictionary with keys "dtype" and "shape".
154+ Args :
155+ metadata: A dictionary with keys "dtype" and "shape".
156156 """
157157 self .group .send_obj (metadata , self .target_rank_for_send )
158158
@@ -161,8 +161,8 @@ def _recv_metadata(self) -> Metadata:
161161 Receive the metadata dictionary from the target rank.
162162
163163 Returns:
164- - metadata: A dictionary with keys "dtype" and "shape" describing
165- the tensor.
164+ metadata: A dictionary with keys "dtype" and "shape"
165+ describing the tensor.
166166 """
167167 return self .group .recv_obj (self .target_rank_for_recv )
168168
@@ -171,9 +171,9 @@ def _send_impl(self, tensor: Optional[torch.Tensor]) -> None:
171171 The actual implementation of sending the tensor and its metadata to the
172172 target rank.
173173
174- Parameters :
175- - tensor: The input tensor to be sent, or None if no tensor is
176- being sent.
174+ Args :
175+ tensor: The input tensor to be sent, or ` None` if no tensor is
176+ being sent.
177177 """
178178 metadata = self ._make_metadata (tensor )
179179 self ._send_metadata (metadata )
@@ -187,7 +187,7 @@ def _recv_impl(self) -> Optional[torch.Tensor]:
187187 the target rank.
188188
189189 Returns:
190- - buffer: The received tensor, or None if no tensor is received.
190+ buffer: The received tensor, or ` None` if no tensor is received.
191191 """
192192 metadata = self ._recv_metadata ()
193193 if metadata ["dtype" ] is None :
@@ -227,8 +227,8 @@ def send_tensor(self, tensor: Optional[torch.Tensor]) -> None:
227227 Sends a tensor and its metadata to the destination rank in a
228228 non-blocking way.
229229
230- Parameters :
231- - tensor: The tensor to send, or None if no tensor is being sent.
230+ Args :
231+ tensor: The tensor to send, or ` None` if no tensor is being sent.
232232 """
233233 if self .transport_thread is None :
234234 self .transport_thread = ThreadPoolExecutor (max_workers = 1 )
@@ -250,8 +250,8 @@ def recv_tensor(self) -> Optional[torch.Tensor]:
250250 """
251251 Receives a tensor and its metadata from the source rank. Blocking call.
252252
253- Returns :
254- - tensor: The received tensor, or None if no tensor is received.
253+ Args :
254+ tensor: The received tensor, or ` None` if no tensor is received.
255255 """
256256 if self .transport_thread is None :
257257 self .transport_thread = ThreadPoolExecutor (max_workers = 1 )
0 commit comments