__init__

Documentation

class treetensor.torch.Size(data, *args, constraint=None, **kwargs)[source]
__init__(data, constraint=None)[source]

In treetensor.torch.Size, it’s similar with the original torch.Size.

Examples:

>>> import torch
>>> import treetensor.numpy as ttorch
>>> ttorch.Size([1, 2, 3])
torch.Size([1, 2, 3])

>>> ttorch.Size({
...     'a': [1, 2, 3],
...     'b': {'x': [3, 4, ]},
...     'c': [5],
... })
<Size 0x7fe00b115970>
├── a --> torch.Size([1, 2, 3])
├── b --> <Size 0x7fe00b115250>
│   └── x --> torch.Size([3, 4])
└── c --> torch.Size([5])