google-deepmind / google-deepmind/tree
`map_structure` seems to be slower than python implementation
- 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
Assessment
This issue has not been assessed yet.