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.6.tar.gz
(5.8 kB
view hashes)
Built Distribution
Close
Hashes for torchstruct-0.1.6-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | e1b1019be87ec2de0387b9909d496039dc5583e8f7f678a85d87658b73ed64ff |
|
MD5 | 539b3a7e996edca0de2084ec14558ad9 |
|
BLAKE2b-256 | 91118cf4a900141fb491f972c397fab960c88bc3f3040ea09c80521ad1c2825f |