No project description provided
Project description
torchstruct
Wrap your multiple torch.Tensor
s into single TensorStruct
and use it like you are using torch.Tensor
.
Installation
pip install torchstruct
Testing
PYTHONPATH=. pytest
Examples
import torch
from torchstruct import TensorStruct
# Initialization
ts = TensorStruct.zeros({
'obs': (2,),
'rew': (1,),
'done': (1,)
}, prefix_shape=(10,), dtype=torch.float32, device='cpu')
raw_data = {
'obs': torch.randn((10, 2)),
'rew': torch.randn((10, 1)),
'done': torch.randn((10, 1))
}
# Assigning
ts[:] = raw_data
# Indexing
ts[2:4]
ts['rew']
# Calling PyTorch methods
ts.unsqueeze(dim=0)
ts.sum(dim=0)
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
torchstruct-0.1.4.tar.gz
(5.7 kB
view hashes)
Built Distribution
Close
Hashes for torchstruct-0.1.4-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | c44db368b752df0c7bf9fe3fadb1689ca6c0691977d63a953b4e9173a7dcab28 |
|
MD5 | 93d0eb884c0e094bad405d13df81edfd |
|
BLAKE2b-256 | a2d9703f9f6bc30fc6e2bb2be1cdbe2adf148e91baf3a7c92b7d6c622e4540e8 |