Test calling `set_current_context(..., require_parent=False).
(self)
| 168 | pass |
| 169 | |
| 170 | def test_nested_context(self): |
| 171 | """Test calling `set_current_context(..., require_parent=False).""" |
| 172 | module1 = _new_module("test1") |
| 173 | module2 = _new_module("test2") |
| 174 | context1 = InvocationContext( |
| 175 | name="context1", |
| 176 | parent=None, # root context |
| 177 | module=module1, |
| 178 | is_training=True, |
| 179 | prng_key=jax.random.PRNGKey(123), |
| 180 | state={"x": 1}, |
| 181 | output_collection=new_output_collection(), |
| 182 | ) |
| 183 | context2 = context1.add_child("context2", module=module2, state={"x": 2}) |
| 184 | context2.parent = None |
| 185 | with set_current_context(context2, require_parent=False) as ctx: |
| 186 | self.assertEqual(ctx.parent, None) |
| 187 | |
| 188 | def test_context_stack_mutlithread(self): |
| 189 | module1 = _new_module("root") |
nothing calls this directly
no test coverage detected