Comparison Between TreeValue and Jax LibTree

In this section, we will take a look at the feature and performance of the jax-libtree library, which is developed by Google.

_TREE_DATA_1 = {'a': 1, 'b': 2, 'x': {'c': 3, 'd': 4}}

Mapping Operation

TreeValue’s Mapping

from treevalue import mapping, FastTreeValue

t = FastTreeValue(_TREE_DATA_1)
mapping(t, lambda x: x ** 2)

<FastTreeValue 0x7f8014317e80>
├── 'a' --> 1
├── 'b' --> 4
└── 'x' --> <FastTreeValue 0x7f8014317dc0>
    ├── 'c' --> 9
    └── 'd' --> 16

%timeit mapping(t, lambda x: x ** 2)

3.65 µs ± 21.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

mapping(t, lambda x, p: (x ** 2, p))

<FastTreeValue 0x7f8014025670>
├── 'a' --> (1, ('a',))
├── 'b' --> (4, ('b',))
└── 'x' --> <FastTreeValue 0x7f80140257c0>
    ├── 'c' --> (9, ('x', 'c'))
    └── 'd' --> (16, ('x', 'd'))

%timeit mapping(t, lambda x, p: (x ** 2, p))

4.04 µs ± 111 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

pytree’s tree_map

from jax.tree_util import tree_map

tree_map(lambda x: x ** 2, _TREE_DATA_1)

{'a': 1, 'b': 4, 'x': {'c': 9, 'd': 16}}

%timeit tree_map(lambda x: x ** 2, _TREE_DATA_1)

6.59 µs ± 41.5 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

Flatten and Unflatten Operation

TreeValue’s Performance

from treevalue import flatten, flatten_keys, flatten_values

t_flatted = flatten(t)
t_flatted

[(('a',), 1), (('b',), 2), (('x', 'c'), 3), (('x', 'd'), 4)]

%timeit flatten(t)

869 ns ± 10.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

from treevalue import flatten_keys

flatten_keys(t)

[('a',), ('b',), ('x', 'c'), ('x', 'd')]

%timeit flatten_keys(t)

732 ns ± 4.65 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

from treevalue import flatten_values

flatten_values(t)

[1, 2, 3, 4]

%timeit flatten_values(t)

540 ns ± 6.72 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

from treevalue import unflatten

unflatten(t_flatted)

<TreeValue 0x7f8014322100>
├── 'a' --> 1
├── 'b' --> 2
└── 'x' --> <TreeValue 0x7f80143222e0>
    ├── 'c' --> 3
    └── 'd' --> 4

%timeit unflatten(t_flatted)

987 ns ± 9.44 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

pytree’s Performance

from jax.tree_util import tree_flatten

leaves, treedef = tree_flatten(_TREE_DATA_1)
print('Leaves:', leaves)
print('Treedef:', treedef)

Leaves: [1, 2, 3, 4]
Treedef: PyTreeDef({'a': *, 'b': *, 'x': {'c': *, 'd': *}})

%timeit tree_flatten(_TREE_DATA_1)

2.36 µs ± 40.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

from jax.tree_util import tree_unflatten

tree_unflatten(treedef, leaves)

{'a': 1, 'b': 2, 'x': {'c': 3, 'd': 4}}

%timeit tree_unflatten(treedef, leaves)

1.02 µs ± 9.66 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

All Operation

TreeValue’s Performance

all(flatten_values(t))

True

%timeit all(flatten_values(t))

688 ns ± 3.52 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

pytree.tree_all’s performance

from jax.tree_util import tree_all

tree_all(_TREE_DATA_1)

True

%timeit tree_all(_TREE_DATA_1)

2.52 µs ± 10.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

Reduce Operation

TreeValue’s Reduce

from functools import reduce

def _flatten_reduce(tree):
    values = flatten_values(tree)
    return reduce(lambda x, y: x + y, values)

_flatten_reduce(t)

10

%timeit _flatten_reduce(t)

1.35 µs ± 5.56 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

def _flatten_reduce_with_init(tree):
    values = flatten_values(tree)
    return reduce(lambda x, y: x + y, values, 0)

_flatten_reduce_with_init(t)

10

%timeit _flatten_reduce_with_init(t)

1.5 µs ± 1.52 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

pytree.tree_reduce

from jax.tree_util import tree_reduce

tree_reduce(lambda x, y: x + y, _TREE_DATA_1)

10

%timeit tree_reduce(lambda x, y: x + y, _TREE_DATA_1)

3.29 µs ± 9.34 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

tree_reduce(lambda x, y: x + y, _TREE_DATA_1, 0)

10

%timeit tree_reduce(lambda x, y: x + y, _TREE_DATA_1, 0)

3.53 µs ± 5.38 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

Structure Transpose

Subside and Rise in TreeValue

from treevalue import subside

value = {
    'a': FastTreeValue({'a': 1, 'b': {'x': 2, 'y': 3}}),
    'b': FastTreeValue({'a': 10, 'b': {'x': 20, 'y': 30}}),
    'c': {
        'x': FastTreeValue({'a': 100, 'b': {'x': 200, 'y': 300}}),
        'y': FastTreeValue({'a': 400, 'b': {'x': 500, 'y': 600}}),
    },
}
subside(value)

<FastTreeValue 0x7f7fdd1e7e20>
├── 'a' --> {'a': 1, 'b': 10, 'c': {'x': 100, 'y': 400}}
└── 'b' --> <FastTreeValue 0x7f7fdd1e70a0>
    ├── 'x' --> {'a': 2, 'b': 20, 'c': {'x': 200, 'y': 500}}
    └── 'y' --> {'a': 3, 'b': 30, 'c': {'x': 300, 'y': 600}}

%timeit subside(value)

16.8 µs ± 54.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

from treevalue import raw, rise

value = FastTreeValue({
    'a': raw({'a': 1, 'b': {'x': 2, 'y': 3}}),
    'b': raw({'a': 10, 'b': {'x': 20, 'y': 30}}),
    'c': {
        'x': raw({'a': 100, 'b': {'x': 200, 'y': 300}}),
        'y': raw({'a': 400, 'b': {'x': 500, 'y': 600}}),
    },
})
rise(value)

{'b': {'y': <FastTreeValue 0x7f7fdd1e7160>
  ├── 'a' --> 3
  ├── 'b' --> 30
  └── 'c' --> <FastTreeValue 0x7f8014025280>
      ├── 'x' --> 300
      └── 'y' --> 600,
  'x': <FastTreeValue 0x7f801402f280>
  ├── 'a' --> 2
  ├── 'b' --> 20
  └── 'c' --> <FastTreeValue 0x7f7fdd1e7070>
      ├── 'x' --> 200
      └── 'y' --> 500},
 'a': <FastTreeValue 0x7f801402fd00>
 ├── 'a' --> 1
 ├── 'b' --> 10
 └── 'c' --> <FastTreeValue 0x7f801402f910>
     ├── 'x' --> 100
     └── 'y' --> 400}

%timeit rise(value)

19.2 µs ± 43.5 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

vt = {'a': None, 'b': {'x': None, 'y': None}}
rise(value, template=vt)

{'a': <FastTreeValue 0x7f7fdd1e7730>
 ├── 'a' --> 1
 ├── 'b' --> 10
 └── 'c' --> <FastTreeValue 0x7f7fdd1e7c70>
     ├── 'x' --> 100
     └── 'y' --> 400,
 'b': {'x': <FastTreeValue 0x7f7fdd1e7040>
  ├── 'a' --> 2
  ├── 'b' --> 20
  └── 'c' --> <FastTreeValue 0x7f7fdd1e7550>
      ├── 'x' --> 200
      └── 'y' --> 500,
  'y': <FastTreeValue 0x7f7fdd1e7a00>
  ├── 'a' --> 3
  ├── 'b' --> 30
  └── 'c' --> <FastTreeValue 0x7f7fdd1e7f70>
      ├── 'x' --> 300
      └── 'y' --> 600}}

%timeit rise(value, template=vt)

15.1 µs ± 34.6 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

pytree.tree_transpose

from jax.tree_util import tree_structure, tree_transpose

sto = tree_structure({'a': 1, 'b': 2, 'c': {'x': 3, 'y': 4}})
sti = tree_structure({'a': 1, 'b': {'x': 2, 'y': 3}})

value = (
    {'a': 1, 'b': {'x': 2, 'y': 3}},
    {
        'a': {'a': 10, 'b': {'x': 20, 'y': 30}},
        'b': [
            {'a': 100, 'b': {'x': 200, 'y': 300}},
            {'a': 400, 'b': {'x': 500, 'y': 600}},
        ],
    }
)
tree_transpose(sto, sti, value)

{'a': {'a': 1, 'b': 10, 'c': {'x': 100, 'y': 400}},
 'b': {'x': {'a': 2, 'b': 20, 'c': {'x': 200, 'y': 500}},
  'y': {'a': 3, 'b': 30, 'c': {'x': 300, 'y': 600}}}}

%timeit tree_transpose(sto, sti, value)

16.6 µs ± 21.1 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)