Function assert_forward_equal

tests/test_lazy/lazy_init_utils.py:41–66

(
    m1: torch.nn.Module,
    m2: torch.nn.Module,
    data_gen_fn: Callable[[], dict],
    output_transform_fn: Callable[[Any], dict],
)

def assert_forward_equal(
    m1: torch.nn.Module,
    m2: torch.nn.Module,
    data_gen_fn: Callable[[], dict],
    output_transform_fn: Callable[[Any], dict],
) -> None:
    data = data_gen_fn()

    m1.eval()
    m2.eval()
    # run forward
    with torch.no_grad():
        outputs1 = m1(**data)
        outputs2 = m2(**data)

    # compare output
    transformed_out1 = output_transform_fn(outputs1)
    transformed_out2 = output_transform_fn(outputs2)

    assert len(transformed_out1) == len(transformed_out2)

    for key, out1 in transformed_out1.items():
        out2 = transformed_out2[key]
        assert torch.allclose(
            out1, out2, atol=1e-5
        ), f"{m1.__class__.__name__} has inconsistent outputs, {out1} vs {out2}"


def check_lazy_init(