ASEM000 / pytreeclass

Visualize, create, and operate on pytrees in the most intuitive way possible.
https://pytreeclass.rtfd.io/en/latest
Apache License 2.0
42 stars 2 forks source link

Faster trace #41

Closed ASEM000 closed 1 year ago

ASEM000 commented 1 year ago

Improved performance in tracing, affects any operation using tree_map_with_trace

New registering format (i.e. no need to wrap it with a LeafTrace object)


def _dict_trace_func(tree: dict) :
    names = (f"['{k}']" for k in tree)
    types = (type(tree[key]) for key in tree)
    index = range(len(tree))
    width = (len(tree) for _ in range(len(tree)))
    metas = ({"repr": not k.startswith("_")} for k in tree)
    return [*zip(names, types, index, width, metas)]