bagua.torch_api.tensor¶
Module Contents¶
- class bagua.torch_api.tensor.BaguaTensor¶
This class patch torch.Tensor with additional methods.
- is_bagua_tensor(self)¶
- Return type
bool
- ensure_bagua_tensor(self, name=None)¶
Convert a PyTorch tensor or parameter to Bagua tensor inplace and return it. A Bagua tensor is required to use Bagua’s communication algorithms.
- Parameters
name (Optional[str]) – the unique name of the tensor
- Returns
The original tensor with Bagua tensor attributes initialized.
- to_bagua_tensor(self, name=None)¶
Create a new Bagua tensor from a PyTorch tensor or parameter and return it. The original tensor is not changed. A Bagua tensor is required to use Bagua’s communication algorithms.
- Parameters
name (Optional[str]) – the unique name of the tensor
- Returns
The new Bagua tensor sharing the same storage with the original tensor.
- bagua_backend_tensor(self)¶
- Returns
The raw Bagua backend tensor.
- Return type
bagua_core.BaguaTensorPy
- bagua_ensure_grad(self)¶
Return the gradient of current parameter. Create a zero gradient tensor if not exist.
- Return type
torch.Tensor
- bagua_mark_communication_ready(self)¶
Mark a Bagua tensor ready for scheduled operations execution.
- bagua_mark_communication_ready_without_synchronization(self)¶
Mark a Bagua tensor ready immediately, without CUDA event synchronization.
- bagua_set_storage(self, storage, storage_offset=0)¶
Sets the underlying storage using an existing torch.Storage.
- Parameters
storage (torch.Storage) – the storage to use
storage_offset (int) – the offset in the storage