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
Built Distribution
Close
Hashes for torchstruct-0.1.5.post0-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 4c134151dac90db31a92e91451da8e7fff3ee0c8162a5dd2cf316fac4ab3b661 |
|
MD5 | 26afe1930a204aa6c5e7c1d3bfbbe499 |
|
BLAKE2b-256 | ecb9a1cd8f003c8c405e0791fb836cdb3309acce4b4819b9e8cf63342b16e32f |