Class TestCase

axlearn/common/test_utils.py:163–328

Base test class.

class TestCase(parameterized.TestCase):
    """Base test class."""

    @property
    def data_dir(self):
        return "FAKE"

    def setUp(self):
        super().setUp()
        push_data_dir(self.data_dir)
        utils_spmd.setup(jax_backend=self._jax_backend())

    def _jax_backend(self) -> str:
        # Setup without distributed initialization by default.
        return "cpu"

    def tearDown(self) -> None:
        super().tearDown()
        self.assertEqual(pop_data_dir(), self.data_dir)

    def _compute_layer_outputs(
        self,
        *,
        test_layer: BaseLayer,
        ref_layer: Any,
        test_inputs: Any,
        ref_inputs: Any,
        parameters_from_ref_layer: ParameterConversionFn,
        require_same_num_params: bool = True,
        require_same_tree_structure: bool = True,
    ) -> tuple[Any, Any]:
        layer_params = test_layer.initialize_parameters_recursively(prng_key=jax.random.PRNGKey(0))
        total_num_layer_params = 0
        for name, param in flatten_items(layer_params):
            logging.info("Test: %s=%s", name, param.shape)
            total_num_layer_params += param.size
        logging.info("Total test layer params: %s", total_num_layer_params)

        total_num_ref_params = 0
        for name, param in ref_layer.named_parameters(recurse=True):
            logging.info("Ref: %s=%s", name, param.shape)
            total_num_ref_params += param.numel()
        for name, param in ref_layer.named_buffers(recurse=True):
            logging.info("Ref buffer: %s=%s", name, param.shape)
        logging.info("Total ref layer params: %s", total_num_ref_params)

        if require_same_num_params:
            self.assertEqual(total_num_layer_params, total_num_ref_params)

        params_from_ref = parameters_from_ref_layer(ref_layer, dst_layer=test_layer)

        # Test that leaves have same paths and shapes. Note that trees with different structures can
        # still have the same leaves: {"a": 1, "b": {}} will match {"a": 1}, because there are no
        # leaves under subtree "b".
        self.assertSequenceEqual(
            [f"{name}={value.shape}" for name, value in flatten_items(params_from_ref)],
            [f"{name}={value.shape}" for name, value in flatten_items(layer_params)],
        )