Class ParametrizeTestModule

tests/test_parametrize.py:20–32

Test module with different parameter shapes for testing parametrization.

class ParametrizeTestModule(nn.Module):
    """Test module with different parameter shapes for testing parametrization."""

    def __init__(self, device="cpu", dtype=torch.float32):
        super().__init__()
        # 2D parameter (typical weight matrix)
        self.weight_2d = nn.Parameter(torch.randn(1024, 1024, device=device, dtype=dtype))
        # 3D parameter (MoE expert weights - the main use case for this feature)
        self.expert_weights = nn.Parameter(torch.randn(8, 512, 256, device=device, dtype=dtype))
        # 1D parameter (bias-like)
        self.bias_1d = nn.Parameter(torch.randn(1024, device=device, dtype=dtype))
        # Non-parameter attribute (should not be quantizable)
        self.not_param = torch.randn(32, device=device, dtype=dtype)


@pytest.mark.parametrize("device", get_available_devices())