opendilab/treevalue

[BUG] When function returns a treevalue

HansBug opened this issue · 0 comments

from treevalue import FastTreeValue


def func(x):
    return FastTreeValue({
        'x': x, 'y': x ** 2,
    })


f = FastTreeValue({
    'x': func,
    'y': {
        'z': func,
    }
})
v = FastTreeValue({'x': 2, 'y': {'z': 34}})
r1 = f(v)
print(r1)

The output is

<FastTreeValue 0x7f13ff6491f0>
├── 'x' --> <FastTreeValue 0x7f13ff6491c0>
│           ├── 'x' --> 2
│           └── 'y' --> 4
└── 'y' --> <FastTreeValue 0x7f13ff649220>
    └── 'z' --> <FastTreeValue 0x7f13ff649130>
                ├── 'x' --> 34
                └── 'y' --> 1156

But the correct output should be

<FastTreeValue 0x7f13ff6498b0>
├── 'x' --> <FastTreeValue 0x7f13ff649a00>
│   ├── 'x' --> 2
│   └── 'y' --> 4
└── 'y' --> <FastTreeValue 0x7f13ff6499a0>
    └── 'z' --> <FastTreeValue 0x7f13ff649910>
        ├── 'x' --> 34
        └── 'y' --> 1156