def test_fabric_module_method_lookup():
    """Test that access to methods warns about improper use when a wrapper from a strategy is involved."""

    class OriginalModule(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.submodule = torch.nn.Linear(2, 3)

        def forward(self, x):
            return x

        def method_without_module_invocation(self):
            return 100

        def method_with_submodule_invocation(self):
            self.submodule(torch.rand(2, 2))
            return 101

        def method_with_self_invocation(self):
            self(None)
            return 102

    class ModuleWrapper(torch.nn.Module):
        def __init__(self, module):
            super().__init__()
            self.wrapped = module

        def forward(self, *args, **kwargs):
            return self.wrapped(*args, **kwargs)

    # Regular case: forward_module == original_module -> no warnings
    original_module = OriginalModule()
    fabric_module = _FabricModule(forward_module=original_module, strategy=Mock(), original_module=original_module)
    assert fabric_module.method_without_module_invocation() == 100

    # Special case: original module wrapped by forward module: -> error if method requires rerouting
    original_module = OriginalModule()
    wrapped_module = ModuleWrapper(original_module)
    fabric_module = _FabricModule(
        forward_module=wrapped_module, strategy=Mock(precision=Precision()), original_module=original_module
    )
    assert fabric_module.method_without_module_invocation() == 100
    with pytest.raises(
        RuntimeError, match=r"You are calling the method `OriginalModule.method_with_submodule_invocation\(\)` from"
    ):
        assert fabric_module.method_with_submodule_invocation() == 101
    with pytest.raises(
        RuntimeError, match=r"You are calling the method `OriginalModule.method_with_self_invocation\(\)` from"
    ):
        assert fabric_module.method_with_self_invocation() == 102

    # No error if explicitly marked as forward method
    fabric_module.mark_forward_method("method_with_self_invocation")
    assert fabric_module.method_with_self_invocation() == 102
Test that access to methods warns about improper use when a wrapper from a strategy is involved.
tests/tests_fabric/test_wrappers.py