Method test_none_prng_key

axlearn/common/module_test.py:235–249

(self)

            self.assertEqual(xs, {1, 2})

    def test_none_prng_key(self):
        module1 = _new_module("test1")
        module2 = _new_module("test2")
        context1 = InvocationContext(
            name="context1",
            parent=None,
            module=module1,
            prng_key=None,
            is_training=True,
            state={"x": 1},
            output_collection=new_output_collection(),
        )
        with set_current_context(context1):
            context2 = context1.add_child("context2", module=module2, state={"x": 2})
            self.assertEqual(context2.prng_key, None)

    class MySummary(summary.Summary):
        val: str