Method test_drop_output

axlearn/common/module_test.py:925–992

(self)

            )

    def test_drop_output(self):
        num_iters = 3
        child_name_prefix = "foo"

        # By default, nothing is dropped.
        with self._dummy_context() as ctx:
            self._invoke(num_iters=num_iters, xs={}, child_name_prefix=child_name_prefix)
            self.assertNestedEqual(
                {
                    "output": jnp.array([[0, 0], [1, 1], [2, 2]]),
                    "with_carry": {
                        "output": jnp.array([[0, 0], [2, 2], [4, 4]]),
                        "with_state": {"output": jnp.array([[0, 10], [2, 12], [4, 14]])},
                    },
                },
                ctx.output_collection.module_outputs[child_name_prefix]["nested"],
            )
            # Summary values are put in `{child_name_prefix}{i}`.
            self.assertNestedEqual(
                {
                    f"{child_name_prefix}{i}": {"carry": WeightedSummary(i, 1)}
                    for i in range(num_iters)
                },
                ctx.output_collection.summaries,
            )

        # Invoke with an output dropper that drops everything.
        with self._dummy_context() as ctx:
            self._invoke(
                num_iters=num_iters,
                xs={},
                drop_output=lambda _: True,
                child_name_prefix=child_name_prefix,
            )
            self.assertNestedEqual({}, ctx.output_collection.module_outputs)

        # Invoke with an output dropper that matches specific paths.
        with self._dummy_context() as ctx:
            self._invoke(
                num_iters=num_iters,
                xs={},
                drop_output=lambda path: match_regex_rules(path, rules=[(".*/with_carry.*", True)]),
                child_name_prefix=child_name_prefix,
            )
            self.assertNestedEqual(
                {
                    "output": jnp.array([[0, 0], [1, 1], [2, 2]]),
                },
                ctx.output_collection.module_outputs[child_name_prefix]["nested"],
            )

        # Invoke with an output dropper that matches specific paths.
        with self._dummy_context() as ctx:
            self._invoke(
                num_iters=num_iters,
                xs={},
                drop_output=lambda path: match_regex_rules(path, rules=[(".*/with_state.*", True)]),
                child_name_prefix=child_name_prefix,