pip install di-treetensor

import builtinsimport osfrom functools import partialimport treetensor.torch as torchprint = partial(builtins.print, sep=os.linesep)if __name__ == '__main__':    # create a tree tensor    t = torch.randn({'a': (2, 3), 'b': {'x': (3, 4)}})    print(t)    print(torch.randn(4, 5))  # create a normal tensor    print()    # structure of tree    print('Structure of tree')    print('t.a:', t.a)  # t.a is a native tensor    print('t.b:', t.b)  # t.b is a tree tensor    print('t.b.x', t.b.x)  # t.b.x is a native tensor    print()    # math calculations    print('Math calculation')    print('t ** 2:', t ** 2)    print('torch.sin(t).cos()', torch.sin(t).cos())    print()    # backward calculation    print('Backward calculation')    t.requires_grad_(True)    t.std().arctan().backward()    print('grad of t:', t.grad)    print()    # native operation    # all the ops can be used as the original usage of `torch`    print('Native operation')    print('torch.sin(t.a)', torch.sin(t.a))  # sin of native tensor

<Tensor 0x7f0dae602760>├── a --> tensor([[-1.2672, -1.5817, -0.3141],│                 [ 1.8107, -0.1023,  0.0940]])└── b --> <Tensor 0x7f0dae602820>    └── x --> tensor([[ 1.2224, -0.3445, -0.9980, -0.4085],                      [ 1.5956,  0.8825, -0.5702, -0.2247],                      [ 0.9235,  0.4538,  0.8775, -0.2642]])tensor([[-0.9559,  0.7684,  0.2682, -0.6419,  0.8637],        [ 0.9526,  0.2927, -0.0591,  1.2804, -0.2455],        [ 0.4699, -0.9998,  0.6324, -0.6885,  1.1488],        [ 0.8920,  0.4401, -0.7785,  0.5931,  0.0435]])Structure of treet.a:tensor([[-1.2672, -1.5817, -0.3141],        [ 1.8107, -0.1023,  0.0940]])t.b:<Tensor 0x7f0dae602820>└── x --> tensor([[ 1.2224, -0.3445, -0.9980, -0.4085],                  [ 1.5956,  0.8825, -0.5702, -0.2247],                  [ 0.9235,  0.4538,  0.8775, -0.2642]])t.b.xtensor([[ 1.2224, -0.3445, -0.9980, -0.4085],        [ 1.5956,  0.8825, -0.5702, -0.2247],        [ 0.9235,  0.4538,  0.8775, -0.2642]])Math calculationt ** 2:<Tensor 0x7f0dae602eb0>├── a --> tensor([[1.6057, 2.5018, 0.0986],│                 [3.2786, 0.0105, 0.0088]])└── b --> <Tensor 0x7f0dae60c040>    └── x --> tensor([[1.4943, 0.1187, 0.9960, 0.1669],                      [2.5458, 0.7789, 0.3252, 0.0505],                      [0.8528, 0.2059, 0.7699, 0.0698]])torch.sin(t).cos()<Tensor 0x7f0dae621910>├── a --> tensor([[0.5782, 0.5404, 0.9527],│                 [0.5642, 0.9948, 0.9956]])└── b --> <Tensor 0x7f0dae6216a0>    └── x --> tensor([[0.5898, 0.9435, 0.6672, 0.9221],                      [0.5406, 0.7163, 0.8578, 0.9753],                      [0.6983, 0.9054, 0.7185, 0.9661]])Backward calculationgrad of t:<Tensor 0x7f0dae60c400>├── a --> tensor([[-0.0435, -0.0535, -0.0131],│                 [ 0.0545, -0.0064, -0.0002]])└── b --> <Tensor 0x7f0dae60cbe0>    └── x --> tensor([[ 0.0357, -0.0141, -0.0349, -0.0162],                      [ 0.0476,  0.0249, -0.0213, -0.0103],                      [ 0.0262,  0.0113,  0.0248, -0.0116]])Native operationtorch.sin(t.a)tensor([[-0.9543, -0.9999, -0.3089],        [ 0.9714, -0.1021,  0.0939]], grad_fn=<SinBackward>)

pip install DI-treetensor[potc]