Similar to jax.tree.map(), but vectorizes `fn` on VDict's.
(fn, tree, *rest)
| 465 | |
| 466 | |
| 467 | def vectorized_tree_map(fn, tree, *rest): |
| 468 | """Similar to jax.tree.map(), but vectorizes `fn` on VDict's.""" |
| 469 | |
| 470 | def vectorized_fn(*nodes): |
| 471 | if isinstance(nodes[0], VDict): |
| 472 | if not jax.tree_util.tree_leaves(nodes[0]): |
| 473 | # This can happen when all VDict values are None and cause issues with jax.vmap. |
| 474 | return nodes[0] |
| 475 | nodes = [dict(**node) for node in nodes] |
| 476 | result = jax.vmap(functools.partial(vectorized_tree_map, fn))(*nodes) |
| 477 | return VDict(**result) |
| 478 | return fn(*nodes) |
| 479 | |
| 480 | return jax.tree.map(vectorized_fn, tree, *rest, is_leaf=lambda t: isinstance(t, VDict)) |
| 481 | |
| 482 | |
| 483 | def expand_vdicts(tree: NestedTensor) -> NestedTensor: |