(self)
| 473 | str(jax.tree_util.tree_structure(tree)), |
| 474 | ) |
| 475 | self.assertLen(jax.tree_util.tree_leaves(tree), 2) |
| 476 | |
| 477 | def test_vectorized_tree_map(self): |
| 478 | tree = VDict(a=jnp.arange(10), b=jnp.arange(7) - 3) |
| 479 | self.assertEqual(VDict(a="a", b="b"), tree_paths(tree)) |
| 480 | self.assertNestedAllClose([("a", tree["a"]), ("b", tree["b"])], flatten_items(tree)) |
| 481 | |
| 482 | # Stack 3 trees together. |
| 483 | stacked_tree = jax.tree.map(lambda *xs: jnp.stack(xs), tree, tree, tree) |
| 484 | self.assertEqual(type(stacked_tree), VDict) |
| 485 | self.assertEqual(VDict(a=(3, 10), b=(3, 7)), jax.tree.map(lambda t: t.shape, stacked_tree)) |
| 486 | |
| 487 | # jax.tree.map() treats VDict similarly to dict. |
| 488 | self.assertEqual(VDict(a=45 * 3, b=0), jax.tree.map(lambda t: t.sum(), stacked_tree)) |
| 489 | # vectorized_tree_map() vectorizes 'fn' on VDict and processes the 3 trees separately. |
| 490 | self.assertNestedAllClose( |
| 491 | VDict(a=jnp.asarray([45, 45, 45]), b=jnp.asarray([0, 0, 0])), |
| 492 | vectorized_tree_map(lambda t: t.sum(), stacked_tree), |
| 493 | ) |
| 494 | |
| 495 | # Nested VDict. |
| 496 | tree2 = VDict(c=stacked_tree) |
| 497 | stacked_tree2 = jax.tree.map(lambda *xs: jnp.stack(xs), tree2, tree2) |
| 498 | self.assertEqual( |
| 499 | VDict(c=VDict(a=(2, 3, 10), b=(2, 3, 7))), |
| 500 | jax.tree.map(lambda t: t.shape, stacked_tree2), |
| 501 | ) |
| 502 | self.assertNestedAllClose( |
| 503 | VDict(c=VDict(a=jnp.full([2, 3], 45), b=jnp.full([2, 3], 0))), |
| 504 | vectorized_tree_map(lambda t: t.sum(), stacked_tree2), |
| 505 | ) |
| 506 |
nothing calls this directly
no test coverage detected