def test_max_norm_parameter(max_norm):
    """
    Test that different max_norm values are respected.
    """
    in_features = 3
    out_features = 2
    layer = LinearWithConstraint(in_features=in_features, out_features=out_features, max_norm=max_norm)
    with torch.no_grad():
        layer.weight.data = torch.tensor([[3.0, 0.0, 0.0],
                                         [0.0, 4.0, 0.0]])
    input = torch.randn(1, in_features)
    layer(input)
    weight_norms = layer.weight.data.norm(p=2, dim=1)
    assert torch.all(weight_norms <= max_norm + 1e-6), (
        f"For max_norm {max_norm}, weight norms {weight_norms} exceed max_norm"
    )

    Test that different max_norm values are respected.
    
test/unit_tests/models/test_modules.py