(self)
| 96 | |
| 97 | class InvocationContextTest(test_utils.TestCase): |
| 98 | def test_set_state_update(self): |
| 99 | for levels in range(4): |
| 100 | context = InvocationContext( |
| 101 | name="root", |
| 102 | parent=None, |
| 103 | module=_new_module("test"), |
| 104 | is_training=True, |
| 105 | prng_key=None, |
| 106 | state={}, |
| 107 | output_collection=new_output_collection(), |
| 108 | ) |
| 109 | descendant = context |
| 110 | for level in range(levels): |
| 111 | name = f"mod{level}" |
| 112 | descendant.module._add_child(name, _Module.default_config()) |
| 113 | descendant = descendant.add_child(name) |
| 114 | descendant.set_state_update((1, 2, 3)) |
| 115 | self.assertEqual(descendant.get_state_updates(), (1, 2, 3)) |
| 116 | |
| 117 | descendant_state_update = context.output_collection.state_updates |
| 118 | for level in range(levels): |
| 119 | name = f"mod{level}" |
| 120 | descendant_state_update = descendant_state_update[name] |
| 121 | self.assertEqual(descendant_state_update, (1, 2, 3), msg=f"levels={levels}") |
| 122 | |
| 123 | def test_context_output_collection(self): |
| 124 | context = InvocationContext( |
nothing calls this directly
no test coverage detected