Method _dummy_context

axlearn/common/module_test.py:882–894

(self)

    @contextlib.contextmanager
    def _dummy_context(self):
        # Yields a dummy context.
        context = InvocationContext(
            name="root",
            parent=None,
            module=_new_module("test"),
            is_training=True,
            prng_key=jax.random.PRNGKey(123),
            state=jnp.arange(2, dtype=jnp.int32) * 10,
            output_collection=new_output_collection(),
        )
        with set_current_context(context):
            yield context

    def test_context(self):
        # Running outside of context should error.