google-deepmind / google-deepmind/tree

`map_structure` seems to be slower than python implementation

Open
#107 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1k
Forks
73
PR merge metrics
No merged PRs in 30d

Description

Implementing a basic recursive python implementation to do the same (or similar) is faster than `map_structure`:
```python
import tree
from time import perf_counter
import collections

# recursive implementation
def nested_map(fn, nest):
if isinstance(nest, list):
return [nested_map(fn, v) for v in nest]
elif isinstance(nest, tuple):
return tuple(nested_map(fn, v) for v in nest)
elif isinstance(nest, collections.abc.Mapping):
return {k: nested_map(fn, v) for k, v in nest.items()}
return fn(nest)

args = [1, 2, [3, 4], {"a": 5}]

s = perf_counter()
tree.map_structure(lambda x: x**2, args)
print(perf_counter() - s) # 9.1e-5

s = perf_counter()
nested_map(lambda x: x**2, args)
print(perf_counter() - s) # 9.6e-6
```
This library looks like a really useful tool but doesn't seem to give the results I would have expected.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.