@@ -118,11 +118,11 @@ def _make_metadata(self, tensor: Optional[torch.Tensor]) -> Metadata:
118
118
"""
119
119
Create the metadata as a dictionary based on the input tensor.
120
120
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.
123
123
124
124
Returns:
125
- - metadata: A dictionary with the following keys:
125
+ metadata: A dictionary with the following keys:
126
126
- "dtype": The data type of the tensor or None.
127
127
- "shape": The shape of the tensor or None.
128
128
"""
@@ -135,13 +135,13 @@ def _prepare_recv_buffer(self, metadata: Metadata) -> torch.Tensor:
135
135
"""
136
136
Create a buffer to receive the tensor based on the provided metadata.
137
137
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.
141
141
142
142
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` .
145
145
"""
146
146
return torch .empty (metadata ["shape" ],
147
147
dtype = metadata ["dtype" ],
@@ -151,8 +151,8 @@ def _send_metadata(self, metadata: Metadata):
151
151
"""
152
152
Send the metadata dictionary to the target rank.
153
153
154
- Parameters :
155
- - metadata: A dictionary with keys "dtype" and "shape".
154
+ Args :
155
+ metadata: A dictionary with keys "dtype" and "shape".
156
156
"""
157
157
self .group .send_obj (metadata , self .target_rank_for_send )
158
158
@@ -161,8 +161,8 @@ def _recv_metadata(self) -> Metadata:
161
161
Receive the metadata dictionary from the target rank.
162
162
163
163
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.
166
166
"""
167
167
return self .group .recv_obj (self .target_rank_for_recv )
168
168
@@ -171,9 +171,9 @@ def _send_impl(self, tensor: Optional[torch.Tensor]) -> None:
171
171
The actual implementation of sending the tensor and its metadata to the
172
172
target rank.
173
173
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.
177
177
"""
178
178
metadata = self ._make_metadata (tensor )
179
179
self ._send_metadata (metadata )
@@ -187,7 +187,7 @@ def _recv_impl(self) -> Optional[torch.Tensor]:
187
187
the target rank.
188
188
189
189
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.
191
191
"""
192
192
metadata = self ._recv_metadata ()
193
193
if metadata ["dtype" ] is None :
@@ -227,8 +227,8 @@ def send_tensor(self, tensor: Optional[torch.Tensor]) -> None:
227
227
Sends a tensor and its metadata to the destination rank in a
228
228
non-blocking way.
229
229
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.
232
232
"""
233
233
if self .transport_thread is None :
234
234
self .transport_thread = ThreadPoolExecutor (max_workers = 1 )
@@ -250,8 +250,8 @@ def recv_tensor(self) -> Optional[torch.Tensor]:
250
250
"""
251
251
Receives a tensor and its metadata from the source rank. Blocking call.
252
252
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.
255
255
"""
256
256
if self .transport_thread is None :
257
257
self .transport_thread = ThreadPoolExecutor (max_workers = 1 )
0 commit comments