Method test_context

axlearn/common/module_test.py:896–899

(self)

            yield context

    def test_context(self):
        # Running outside of context should error.
        with self.assertRaisesRegex(ValueError, "Expected current_context()"):
            self._invoke(num_iters=3, xs={})

    def test_basic(self):
        num_iters = 3