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.5.tar.gz
(5.8 kB
view hashes)
Built Distribution
Close
Hashes for torchstruct-0.1.5-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 8613a609ab87b43321b674c626f44f0de4b98d387101a60caec7482013ff8413 |
|
MD5 | 7ad9c9cb2e68acf954b258225522a8ee |
|
BLAKE2b-256 | 9f0b770b3d3efe48cc28febe5580f6519cef52620c63f1165c7fa445cce53792 |