Repository: ray-project/ray
File: rllib/core/models/tests/test_encoder_layernorm.py (Python)

Code before:
from ray.rllib.core.models.configs import MLPEncoderConfig
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
from ray.rllib.utils.framework import try_import_torch

torch, nn = try_import_torch()


def count_layernorm_modules(module):
    """Count LayerNorm layers in a module and its submodules."""
    return sum(1 for m in module.modules() if isinstance(m, nn.LayerNorm))


def get_layernorm_positions(sequential_net):
    """Get indices and normalized shapes of LayerNorm layers in a Sequential net."""
    result = []
    for i, layer in enumerate(sequential_net):
        if isinstance(layer, nn.LayerNorm):
            result.append((i, tuple(layer.normalized_shape)))
    return result


class TestEncoderLayerNormConfig(unittest.TestCase):
    """Tests for LayerNorm in MLPEncoderConfig with various config combinations."""

    def test_hidden_only_layernorm(self):
        """With hidden_layer_use_layernorm=True, output_layer_use_layernorm=False."""
        config = MLPEncoderConfig(
            input_dims=[10],
            hidden_layer_dims=[16, 16],
            hidden_layer_activation="relu",
            hidden_layer_use_layernorm=True,
            output_layer_use_layernorm=False,
            output_layer_dim=16,
            output_layer_activation="relu",
        )
        encoder = config.build(framework="torch")
        layernorm_count = count_layernorm_modules(encoder.net)
        self.assertEqual(
            layernorm_count, 2, "Expected 2 LayerNorms after hidden layers"
        )
        # No LayerNorm after output layer
        positions = get_layernorm_positions(encoder.net.mlp)
        self.assertEqual(len(positions), 2)
        self.assertEqual(positions[0][1], (16,))
        self.assertEqual(positions[1][1], (16,))

    def test_output_only_layernorm(self):
        """With hidden_layer_use_layernorm=False, output_layer_use_layernorm=True."""
        config = MLPEncoderConfig(
            input_dims=[10],

Review comment (refactor):
All of these tests look like the same test with different configs. If correct, update to use `pytest.mark.parametrize` to simplify the testing

Code after:
from ray.rllib.core.models.configs import MLPEncoderConfig
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
from ray.rllib.utils.framework import try_import_torch

torch, nn = try_import_torch()


def count_layernorm_modules(module):
    """Count LayerNorm layers in a module and its submodules."""
    return sum(1 for m in module.modules() if isinstance(m, nn.LayerNorm))


def get_layernorm_positions(sequential_net):
    """Get indices and normalized shapes of LayerNorm layers in a Sequential net."""
    result = []
    for i, layer in enumerate(sequential_net):
        if isinstance(layer, nn.LayerNorm):
            result.append((i, tuple(layer.normalized_shape)))
    return result


def get_full_model_config(overrides):
    """Merge overrides with DefaultModelConfig for Catalog compatibility."""
    return dataclasses.asdict(DefaultModelConfig()) | overrides


# ---------------------------------------------------------------------------
# MLPEncoderConfig parametrized tests
# ---------------------------------------------------------------------------

MLP_ENCODER_CONFIGS = [
    pytest.param(
        {
            "hidden_layer_use_layernorm": True,
            "output_layer_use_layernorm": False,
            "hidden_layer_dims": [16, 16],
            "output_layer_dim": 16,
            "expected_count": 2,
            "expected_positions": [(16,), (16,)],
        },
        id="hidden_only_layernorm",
    ),
    pytest.param(
        {
            "hidden_layer_use_layernorm": False,
            "output_layer_use_layernorm": True,
            "hidden_layer_dims": [16],
            "output_layer_dim": 16,
            "expected_count": 1,
            "expected_positions": [(16,)],