"""Tests for toy models from example_models.py.

All existing validation + visualization tests migrated from test_validation_and_visuals.py,
plus new API coverage tests.
"""

from os.path import join as opj

import pytest
import torch

from conftest import VIS_OUTPUT_DIR

import example_models
from torchlens import (
    log_forward_pass,
    get_model_metadata,
    show_model_graph,
    validate_forward_pass,
)


# =============================================================================
# Simple operations
# =============================================================================


@pytest.mark.smoke
def test_model_simple_ff(default_input1):
    model = example_models.SimpleFF()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "simple_ff"),
    )


def test_model_inplace_funcs(default_input1):
    model = example_models.InPlaceFuncs()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "inplace_funcs"),
    )


def test_model_simple_internally_generated(default_input1):
    model = example_models.SimpleInternallyGenerated()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "simple_internally_generated"),
    )


def test_model_new_tensor_inside(default_input1):
    model = example_models.NewTensorInside()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "new_tensor_inside"),
    )


def test_model_new_tensor_from_numpy(default_input1):
    model = example_models.TensorFromNumpy()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "new_tensor_from_numpy"),
    )


def test_model_simple_random(default_input1):
    model = example_models.SimpleRandom()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "simple_random"),
    )


# =============================================================================
# Dropout
# =============================================================================


def test_dropout_model_real_train(default_input1):
    model = example_models.DropoutModelReal()
    model.train()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "dropout_real_train"),
    )


def test_dropout_model_real_eval(default_input1):
    model = example_models.DropoutModelReal()
    model.eval()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "dropout_real_eval"),
    )


def test_dropout_model_dummy_zero_train(default_input1):
    model = example_models.DropoutModelDummyZero()
    model.train()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "dropout_dummyzero_train"),
    )


def test_dropout_model_dummy_zero_eval(default_input1):
    model = example_models.DropoutModelDummyZero()
    model.eval()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "dropout_dummyzero_eval"),
    )


# =============================================================================
# BatchNorm
# =============================================================================


def test_batchnorm_train(default_input1):
    model = example_models.BatchNormModel()
    model.train()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_buffer_layers=True,
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "batchnorm_train_showbuffer"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_buffer_layers=False,
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "batchnorm_train_invisbuffer"),
    )


def test_batchnorm_eval(default_input1):
    model = example_models.BatchNormModel()
    model.eval()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "batchnorm_eval"),
    )


# =============================================================================
# Tensor operations
# =============================================================================


def test_concat_tensors(default_input1):
    model = example_models.ConcatTensors()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "concat_tensors"),
    )


def test_split_tensor(default_input1):
    model = example_models.SplitTensor()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "split_tensors"),
    )


def test_identity_model(default_input1):
    model = example_models.IdentityModule()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "identity_model"),
    )


def test_assign_tensor(input_2d):
    model = example_models.AssignTensor()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "assigntensor"),
    )


def test_get_and_set_item(default_input1):
    model = example_models.GetAndSetItem()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "get_set_item"),
    )


def test_getitem_tracking(input_2d):
    model = example_models.GetItemTracking()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "getitem_tracking"),
    )


def test_inplace_zero_tensor(default_input1):
    model = example_models.InPlaceZeroTensor()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "inplace_zerotensor"),
    )


def test_slice_operations(default_input1):
    model = example_models.SliceOperations()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "slice_operations"),
    )


def test_dummy_operations(default_input1):
    model = example_models.DummyOperations()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "dummy_operations"),
    )


def test_sametensor_arg(default_input1):
    model = example_models.SameTensorArg()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "same_tensor_arg"),
    )


# =============================================================================
# Multiple inputs / outputs
# =============================================================================


def test_multiple_inputs_arg(default_input1, default_input2, default_input3):
    model = example_models.MultiInputs()
    assert validate_forward_pass(model, [default_input1, default_input2, default_input3])
    show_model_graph(
        model,
        [default_input1, default_input2, default_input3],
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "multiple_inputs_arg"),
    )


def test_multiple_inputs_kwarg(default_input1, default_input2, default_input3):
    model = example_models.MultiInputs()
    assert validate_forward_pass(
        model, [], {"x": default_input1, "y": default_input2, "z": default_input3}
    )
    show_model_graph(
        model,
        [],
        {"x": default_input1, "y": default_input2, "z": default_input3},
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "multiple_inputs_kwarg"),
    )


def test_multiple_inputs_arg_kwarg_mix(default_input1, default_input2, default_input3):
    model = example_models.MultiInputs()
    assert validate_forward_pass(
        model, [default_input1], {"y": default_input2, "z": default_input3}
    )
    show_model_graph(
        model,
        [],
        {"x": default_input1, "y": default_input2, "z": default_input3},
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "multiple_inputs_arg_kwarg_mix"),
    )


def test_list_input(default_input1, default_input2, default_input3):
    model = example_models.ListInput()
    assert validate_forward_pass(model, [[default_input1, default_input2, default_input3]])
    show_model_graph(
        model,
        [[default_input1, default_input2, default_input3]],
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "list_inputs"),
    )


def test_dict_input(default_input1, default_input2, default_input3):
    model = example_models.DictInput()
    assert validate_forward_pass(
        model, [{"x": default_input1, "y": default_input2, "z": default_input3}]
    )
    show_model_graph(
        model,
        [{"x": default_input1, "y": default_input2, "z": default_input3}],
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "dict_inputs"),
    )


def test_nested_input(default_input1, default_input2, default_input3, default_input4):
    model = example_models.NestedInput()
    assert validate_forward_pass(
        model,
        [
            {
                "list1": [default_input1, default_input2],
                "list2": [default_input3, default_input4],
            }
        ],
    )
    show_model_graph(
        model,
        [
            {
                "list1": [default_input1, default_input2],
                "list2": [default_input3, default_input4],
            }
        ],
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_inputs"),
    )


def test_multi_outputs(default_input1):
    model = example_models.MultiOutputs()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "multi_outputs"),
    )


def test_list_output(default_input1):
    model = example_models.ListOutput()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "list_output"),
    )


def test_dict_output(default_input1):
    model = example_models.DictOutput()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "dict_output"),
    )


def test_nested_output(default_input1):
    model = example_models.NestedOutput()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_output"),
    )


# =============================================================================
# Buffers
# =============================================================================


def test_buffer_model():
    model = example_models.BufferModel()
    model_input = torch.rand(12, 12)
    assert validate_forward_pass(model, model_input)
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "buffer_visible"),
        vis_buffer_layers=True,
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "buffer_invisible"),
        vis_buffer_layers=False,
    )


def test_buffer_rewrite_model():
    model = example_models.BufferRewriteModel()
    model_input = torch.rand(12, 12)
    assert validate_forward_pass(model, model_input)
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_nesting_depth=1,
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "buffer_rewrite_model_visible_unnested_unrolled",
        ),
        vis_buffer_layers=True,
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_nesting_depth=1,
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "buffer_rewrite_model_invisible_unnested_unrolled",
        ),
        vis_buffer_layers=False,
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "buffer_rewrite_model_visible_nested_unrolled",
        ),
        vis_buffer_layers=True,
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "buffer_rewrite_model_invisible_nested_unrolled",
        ),
        vis_buffer_layers=False,
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="rolled",
        vis_nesting_depth=1,
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "buffer_rewrite_model_visible_unnested_rolled",
        ),
        vis_buffer_layers=True,
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="rolled",
        vis_nesting_depth=1,
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "buffer_rewrite_model_invisible_unnested_rolled",
        ),
        vis_buffer_layers=False,
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "buffer_rewrite_model_visible_nested_rolled",
        ),
        vis_buffer_layers=True,
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "buffer_rewrite_model_invisible_nested_rolled",
        ),
        vis_buffer_layers=False,
    )


# =============================================================================
# Branching
# =============================================================================


def test_simple_branching(default_input1):
    model = example_models.SimpleBranching()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "simple_branching"),
    )


def test_conditional_branching(zeros_input, ones_input):
    model = example_models.ConditionalBranching()
    assert validate_forward_pass(model, -ones_input)
    assert validate_forward_pass(model, ones_input)
    show_model_graph(
        model,
        -ones_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "conditional_branching_negative"),
    )
    show_model_graph(
        model,
        ones_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "conditional_branching_positive"),
    )


# =============================================================================
# Nesting
# =============================================================================


def test_repeated_module(default_input1):
    model = example_models.RepeatedModule()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "repeated_module"),
    )


def test_nested_modules(default_input1):
    model = example_models.NestedModules()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_modules_fulldepth"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_nesting_depth=1,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_modules_depth1"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_nesting_depth=2,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_modules_depth2"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_nesting_depth=3,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_modules_depth3"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_nesting_depth=4,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_modules_depth4"),
    )


# =============================================================================
# Loops
# =============================================================================


def test_orphan_tensors(default_input1):
    model = example_models.OrphanTensors()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "orphan_tensors"),
    )


def test_simple_loop_no_param(default_input1):
    model = example_models.SimpleLoopNoParam()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "simple_loop_no_param_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "simple_loop_no_param_rolled"),
    )


def test_same_op_repeat(vector_input):
    model = example_models.SameOpRepeat()
    assert validate_forward_pass(model, vector_input)
    show_model_graph(
        model,
        vector_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "same_op_repeat_unrolled"),
    )
    show_model_graph(
        model,
        vector_input,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "same_op_repeat_rolled"),
    )


def test_repeated_op_type_in_loop(default_input1):
    model = example_models.RepeatedOpTypeInLoop()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "same_op_type_in_loop_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "same_op_type_in_loop_rolled"),
    )


def test_varying_loop_noparam1(default_input1):
    model = example_models.VaryingLoopNoParam1()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "varying_loop_noparam1_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "varying_loop_noparam1_rolled"),
    )


def test_varying_loop_noparam2(default_input1):
    model = example_models.VaryingLoopNoParam2()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "varying_loop_noparam2_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "varying_loop_noparam2_rolled"),
    )


def test_varying_loop_withparam(vector_input):
    model = example_models.VaryingLoopWithParam()
    assert validate_forward_pass(model, vector_input)
    show_model_graph(
        model,
        vector_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "varying_loop_withparam_unrolled"),
    )
    show_model_graph(
        model,
        vector_input,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "varying_loop_withparam_rolled"),
    )


def test_looping_internal_funcs(default_input1):
    model = example_models.LoopingInternalFuncs()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "looping_internal_funcs_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "looping_internal_funcs_rolled"),
    )


def test_looping_from_inputs1(default_input1, default_input2, default_input3):
    model = example_models.LoopingFromInputs1()
    assert validate_forward_pass(model, [default_input1, default_input2, default_input3])
    show_model_graph(
        model,
        [default_input1, default_input2, default_input3],
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "looping_from_inputs1_unrolled"),
    )
    show_model_graph(
        model,
        [default_input1, default_input2, default_input3],
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "looping_from_inputs1_rolled"),
    )


def test_looping_from_inputs2(default_input1, default_input2, default_input3):
    model = example_models.LoopingFromInputs2()
    assert validate_forward_pass(model, [[default_input1, default_input2, default_input3]])
    show_model_graph(
        model,
        [[default_input1, default_input2, default_input3]],
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "looping_from_inputs2_unrolled"),
    )
    show_model_graph(
        model,
        [[default_input1, default_input2, default_input3]],
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "looping_from_inputs2_rolled"),
    )


def test_looping_inputs_and_outputs(default_input1, default_input2, default_input3):
    model = example_models.LoopingInputsAndOutputs()
    assert validate_forward_pass(model, [[default_input1, default_input2, default_input3]])
    show_model_graph(
        model,
        [[default_input1, default_input2, default_input3]],
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "looping_inputs_and_outputs_unrolled",
        ),
    )
    show_model_graph(
        model,
        [[default_input1, default_input2, default_input3]],
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "looping_inputs_and_outputs_rolled"),
    )


def test_stochastic_loop():
    model = example_models.StochasticLoop()
    model_input = torch.full((2, 2), 98.0)
    assert validate_forward_pass(model, model_input)
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "stochastic_loop_unrolled1"),
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "stochastic_loop_rolled1"),
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "stochastic_loop_unrolled2"),
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "stochastic_loop_rolled2"),
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "stochastic_loop_unrolled3"),
    )
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "stochastic_loop_rolled3"),
    )


# =============================================================================
# Recurrent
# =============================================================================


def test_recurrent_params_simple(input_2d):
    model = example_models.RecurrentParamsSimple()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "recurrent_params_simple_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "recurrent_params_simple_rolled"),
    )


def test_recurrent_params_complex(input_2d):
    model = example_models.RecurrentParamsComplex()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "recurrent_params_complex_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "recurrent_params_complex_rolled"),
    )


def test_looping_params_doublenested(input_2d):
    model = example_models.LoopingParamsDoubleNested()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "looping_params_doublenested_unrolled",
        ),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(
            VIS_OUTPUT_DIR,
            "toy-networks",
            "looping_params_doublenested_rolled",
        ),
    )


# =============================================================================
# Module clashes
# =============================================================================


def test_module_looping_clash1(default_input1):
    model = example_models.ModuleLoopingClash1()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_looping_clash1_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_looping_clash1_rolled"),
    )


def test_module_looping_clash2(default_input1):
    model = example_models.ModuleLoopingClash2()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_looping_clash2_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_looping_clash2_rolled"),
    )


def test_module_looping_clash3(default_input1):
    model = example_models.ModuleLoopingClash3()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_looping_clash3_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_looping_clash3_rolled"),
    )


def test_nested_param_free_loops(default_input1):
    """Tests nested loop topology where inner ops have the same equivalence type
    across levels but surrounding ops differ per level.

    4 outer iterations x 3 inner levels = 12 sin operations, all with the same
    equivalence type. They should be ONE group of 12 passes.
    """
    model = example_models.NestedParamFreeLoops()
    assert validate_forward_pass(model, default_input1)
    mh = log_forward_pass(model, default_input1)

    # All sin ops should be in one layer group with 12 passes
    sin_pass_counts = set()
    for label in mh.layer_labels:
        entry = mh[label]
        if entry.func_applied_name == "sin":
            sin_pass_counts.add(entry.layer_passes_total)
    assert sin_pass_counts == {12}, (
        f"sin ops fragmented into groups with pass counts {sin_pass_counts}, expected {{12}}"
    )

    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_param_free_loops_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_param_free_loops_rolled"),
    )


def test_parallel_loops(default_input1):
    model = example_models.ParallelLoops()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "parallel_loops_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "parallel_loops_rolled"),
    )


def test_shared_param_loop_external(input_2d):
    model = example_models.SharedParamLoopExternal()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "shared_param_loop_external_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "shared_param_loop_external_rolled"),
    )


def test_interleaved_shared_param_loops(input_2d):
    model = example_models.InterleavedSharedParamLoops()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "interleaved_shared_param_loops_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "interleaved_shared_param_loops_rolled"),
    )


def test_nested_loops_independent_params(input_2d):
    model = example_models.NestedLoopsIndependentParams()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_loops_independent_params_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_loops_independent_params_rolled"),
    )


def test_self_feeding_no_param(default_input1):
    model = example_models.SelfFeedingNoParam()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "self_feeding_no_param_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "self_feeding_no_param_rolled"),
    )


def test_diamond_loop(input_2d):
    model = example_models.DiamondLoop()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "diamond_loop_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "diamond_loop_rolled"),
    )


def test_accumulator_loop(input_2d):
    model = example_models.AccumulatorLoop()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "accumulator_loop_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "accumulator_loop_rolled"),
    )


def test_single_iteration_loop(input_2d):
    model = example_models.SingleIterationLoop()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "single_iteration_loop_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "single_iteration_loop_rolled"),
    )


def test_long_loop(input_2d):
    model = example_models.LongLoop()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "long_loop_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "long_loop_rolled"),
    )


def test_data_dependent_branch_loop(input_2d):
    model = example_models.DataDependentBranchLoop()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "data_dependent_branch_loop_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "data_dependent_branch_loop_rolled"),
    )


def test_sequential_param_free_loops(default_input1):
    model = example_models.SequentialParamFreeLoops()
    assert validate_forward_pass(model, default_input1)
    mh = log_forward_pass(model, default_input1)
    # Verify the two sequential loops produce SEPARATE groups (not merged)
    sin_groups = set()
    for label in mh.layer_labels:
        entry = mh[label]
        if entry.func_applied_name == "sin":
            sin_groups.add(entry.layer_passes_total)
    assert sin_groups == {3}, f"Sequential sin loops should each have 3 passes, got {sin_groups}"
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "sequential_param_free_loops_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "sequential_param_free_loops_rolled"),
    )


# =============================================================================
# Complex models
# =============================================================================


def test_propertymodel(input_complex):
    model = example_models.PropertyModel()
    assert validate_forward_pass(model, input_complex)
    show_model_graph(
        model,
        input_complex,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "propertymodel"),
    )


def test_ubermodel1(input_2d):
    model = example_models.UberModel1()
    assert validate_forward_pass(model, [[input_2d, input_2d * 2, input_2d * 3]])
    show_model_graph(
        model,
        [[input_2d, input_2d * 2, input_2d * 3]],
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel1"),
    )


def test_ubermodel2():
    model = example_models.UberModel2()
    model_input = torch.rand(2, 1, 3, 3)
    assert validate_forward_pass(model, model_input)
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel2"),
    )


def test_ubermodel3(input_2d):
    model = example_models.UberModel3()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel3_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel3_rolled"),
    )


def test_ubermodel4(input_2d):
    model = example_models.UberModel4()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel4_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel4_rolled"),
    )


def test_ubermodel5():
    model = example_models.UberModel5()
    model_input = torch.rand(1, 1, 3, 3)
    assert validate_forward_pass(model, model_input)
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel5"),
    )


def test_ubermodel6(default_input1):
    model = example_models.UberModel6()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel6_unrolled"),
    )
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel6_rolled"),
    )


def test_ubermodel7(input_2d):
    model = example_models.UberModel7()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel7_unrolled"),
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel7_rolled"),
    )


def test_ubermodel8():
    model = example_models.UberModel8()
    model_input = torch.rand(2, 1, 3, 3)
    assert validate_forward_pass(model, model_input)
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel8"),
    )


def test_ubermodel9():
    model = example_models.UberModel9()
    model_input = torch.rand(1, 1, 3, 3)
    assert validate_forward_pass(model, model_input)
    show_model_graph(
        model,
        model_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "ubermodel9"),
    )


# =============================================================================
# NEW: GeluModel (exists in example_models but had no test)
# =============================================================================


def test_gelu_model(default_input1):
    model = example_models.GeluModel()
    assert validate_forward_pass(model, default_input1)
    show_model_graph(
        model,
        default_input1,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "gelu_model"),
    )


# =============================================================================
# NEW: API coverage tests
# =============================================================================


def test_log_forward_pass_layers_to_save(default_input1):
    """Test layers_to_save parameter selectively saves activations."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1, layers_to_save="all")
    all_labels = mh.layer_labels
    assert len(all_labels) > 0

    # Save only the first layer
    mh_partial = log_forward_pass(model, default_input1, layers_to_save=[all_labels[0]])
    # The layer list should still track all layers
    assert len(mh_partial.layer_labels) > 0
    # But only the requested layer should have saved activations
    saved_entry = mh_partial[all_labels[0]]
    assert saved_entry.tensor_contents is not None


def test_log_forward_pass_save_function_args(default_input1):
    """Test save_function_args=True populates creation_args on entries."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1, save_function_args=True)
    assert mh.save_function_args is True
    # At least one non-input layer should have creation_args
    found = False
    for label in mh.layer_labels:
        entry = mh[label]
        if not entry.is_input_layer and entry.creation_args is not None:
            found = True
            break
    assert found, "No non-input layer had creation_args populated"


def test_log_forward_pass_activation_postfunc(default_input1):
    """Test activation_postfunc applies to saved tensors."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1, activation_postfunc=torch.mean)
    # All saved tensors should be scalar (mean reduces to scalar)
    for label in mh.layer_labels:
        entry = mh[label]
        if entry.tensor_contents is not None:
            assert entry.tensor_contents.dim() == 0, (
                f"Layer {label} tensor should be scalar after torch.mean postfunc"
            )


def test_log_forward_pass_mark_distances(default_input1):
    """Test mark_input_output_distances=True populates distance fields."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1, mark_input_output_distances=True)
    for label in mh.layer_labels:
        entry = mh[label]
        assert entry.min_distance_from_input is not None
        assert entry.max_distance_from_input is not None
        assert entry.min_distance_from_output is not None
        assert entry.max_distance_from_output is not None


def test_get_model_metadata(default_input1):
    """Test get_model_metadata returns ModelLog without saving activations."""
    model = example_models.SimpleFF()
    mh = get_model_metadata(model, default_input1)
    assert len(mh.layer_labels) > 0
    assert mh.num_tensors_total > 0


def test_model_log_getitem_by_index(default_input1):
    """Test ModelLog supports integer indexing."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1)
    first = mh[0]
    last = mh[-1]
    assert first.layer_label == mh.layer_labels[0]
    assert last.layer_label == mh.layer_labels[-1]


def test_model_log_getitem_by_label(default_input1):
    """Test ModelLog supports string label access."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1)
    label = mh.layer_labels[0]
    entry = mh[label]
    assert entry.layer_label == label


def test_model_log_len(default_input1):
    """Test len(ModelLog) matches layer count."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1)
    assert len(mh) == len(mh.layer_labels)


def test_model_log_iter(default_input1):
    """Test iteration over ModelLog yields entries."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1)
    items = list(mh)
    assert len(items) == len(mh.layer_labels)


def test_model_log_layer_labels(default_input1):
    """Test layer_labels, layer_labels_no_pass, layer_labels_w_pass properties."""
    model = example_models.SimpleFF()
    mh = log_forward_pass(model, default_input1)
    assert isinstance(mh.layer_labels, list)
    assert all(isinstance(lbl, str) for lbl in mh.layer_labels)
    assert isinstance(mh.layer_labels_no_pass, list)
    assert all(isinstance(lbl, str) for lbl in mh.layer_labels_no_pass)
    assert isinstance(mh.layer_labels_w_pass, list)
    assert all(isinstance(lbl, str) for lbl in mh.layer_labels_w_pass)


def test_rolled_vs_unrolled_visualization(input_2d):
    """Test both rolled and unrolled modes produce output on recurrent model."""
    import os

    model = example_models.RecurrentParamsSimple()
    unrolled_path = opj(VIS_OUTPUT_DIR, "toy-networks", "test_rolled_vs_unrolled_unrolled")
    rolled_path = opj(VIS_OUTPUT_DIR, "toy-networks", "test_rolled_vs_unrolled_rolled")
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=unrolled_path,
    )
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="rolled",
        vis_outpath=rolled_path,
    )
    # Both should produce output files
    assert os.path.exists(unrolled_path + ".pdf") or os.path.exists(unrolled_path)
    assert os.path.exists(rolled_path + ".pdf") or os.path.exists(rolled_path)


# =============================================================================
# View mutation / child tensor variation tests
# =============================================================================


def test_view_mutation_unsqueeze(input_2d):
    """Mutation through unsqueeze view should be logged without error."""
    model = example_models.ViewMutationUnsqueeze()
    assert validate_forward_pass(model, input_2d)
    mh = log_forward_pass(model, input_2d, save_function_args=True)
    assert mh is not None
    assert len(mh.layer_labels) > 0


def test_view_mutation_reshape(input_2d):
    """Mutation through reshape view should be logged without error."""
    model = example_models.ViewMutationReshape()
    assert validate_forward_pass(model, input_2d)
    mh = log_forward_pass(model, input_2d, save_function_args=True)
    assert mh is not None
    assert len(mh.layer_labels) > 0


def test_view_mutation_transpose(input_2d):
    """Mutation through transpose view should be logged without error."""
    model = example_models.ViewMutationTranspose()
    assert validate_forward_pass(model, input_2d)
    mh = log_forward_pass(model, input_2d, save_function_args=True)
    assert mh is not None
    assert len(mh.layer_labels) > 0


def test_multiple_view_mutations(input_2d):
    """Multiple views mutated independently should be logged without error."""
    model = example_models.MultipleViewMutations()
    assert validate_forward_pass(model, input_2d)
    mh = log_forward_pass(model, input_2d, save_function_args=True)
    assert mh is not None
    assert len(mh.layer_labels) > 0


def test_chained_view_mutation(input_2d):
    """Mutation through chained views should be logged without error."""
    model = example_models.ChainedViewMutation()
    assert validate_forward_pass(model, input_2d)
    mh = log_forward_pass(model, input_2d, save_function_args=True)
    assert mh is not None
    assert len(mh.layer_labels) > 0


def test_output_matches_parent_no_false_positive(input_2d):
    """No mutation model: verify no false-positive child tensor variations."""
    model = example_models.OutputMatchesParent()
    assert validate_forward_pass(model, input_2d)
    mh = log_forward_pass(model, input_2d, save_function_args=True)
    assert mh is not None
    assert len(mh.layer_labels) > 0
    # No layer should have child tensor variations since nothing is mutated
    for label in mh.layer_labels:
        entry = mh[label]
        assert not entry.has_child_tensor_variations, (
            f"Layer {label} should not have child tensor variations in a "
            f"mutation-free model, but has_child_tensor_variations={entry.has_child_tensor_variations}"
        )


# =============================================================================
# Regression: issue #18 — deepcopy hangs on complex tensor wrappers
# =============================================================================


@pytest.mark.timeout(30)
def test_wrapped_input_no_deepcopy_hang():
    """Ensure torchlens handles non-deepcopy-safe input args (issue #18).

    copy.deepcopy would loop infinitely on objects with circular references
    (e.g. ESCNN's GeometricTensor).  The safe_copy_args helper clones tensors
    and leaves other objects as-is, avoiding the hang.
    """
    original_tensor = torch.rand(5, 5)
    original_data = original_tensor.clone()
    wrapper = example_models.TensorWrapper(original_tensor)
    model = example_models.WrappedInputModel()
    mh = log_forward_pass(model, wrapper)
    assert mh is not None
    assert len(mh.layer_labels) > 0
    # Original tensor should not have been mutated (same-device case)
    assert torch.equal(wrapper.tensor, original_data)


# =============================================================================
# Regression: issue #43 — tuple of tensors as single forward arg
# =============================================================================


def test_tuple_input_single_arg():
    """Ensure a tuple of tensors passed as a single arg is not unpacked (issue #43).

    When a model's forward() expects a single argument that IS a tuple,
    torchlens should not split it into multiple positional args.
    """
    model = example_models.TupleInputModel()
    input_tuple = (torch.rand(5, 5), torch.rand(5, 5))
    mh = log_forward_pass(model, input_tuple)
    assert mh is not None
    assert len(mh.layer_labels) > 0
    assert validate_forward_pass(model, input_tuple)


# =============================================================================
# Regression: issue #48 — functional ops at module end rendered as boxes
# =============================================================================


def test_functional_after_submodule_not_box():
    """Functional ops at the end of container modules should be ovals, not boxes (issue #48).

    A torch.relu after a nn.Linear inside a container module should not get
    box-shaped rendering (which is reserved for module outputs).
    """
    from torchlens.visualization.rendering import _get_node_address_shape_color

    model = example_models.FunctionalAfterSubmodule()
    x = torch.rand(2, 5)
    mh = log_forward_pass(model, x)
    # Find the relu layer
    relu_layers = [label for label in mh.layer_labels if "relu" in label.lower()]
    assert len(relu_layers) > 0, "No relu layer found"
    relu_entry = mh[relu_layers[0]]
    _, shape, _ = _get_node_address_shape_color(mh, relu_entry, show_buffer_layers=False)
    assert shape == "oval", f"Functional relu should be oval, got {shape}"


# =============================================================================
# Regression: issue #46 — output layer tensor_contents None with layers_to_save
# =============================================================================


def test_output_layer_saved_with_layers_to_save():
    """Output layers should have tensor_contents even when layers_to_save is a subset (issue #46).

    The output layer copies tensor_contents from its parent, so the parent must
    also be saved during the fast pass.
    """
    model = example_models.FunctionalAfterSubmodule()
    x = torch.rand(2, 5)
    # Save only the relu layer (not explicitly the output)
    mh = log_forward_pass(model, x, layers_to_save=["relu_1"])
    for label in mh.output_layers:
        entry = mh[label]
        assert entry.tensor_contents is not None, (
            f"Output layer {label} should have tensor_contents"
        )


# =============================================================================
# Regression: issue #58 — two-pass fails for stochastic depth models
# =============================================================================


def test_stochastic_depth_layers_to_save():
    """Two-pass with layers_to_save should reproduce the same graph for
    stochastic depth models by restoring the RNG state (issue #58).
    """
    model = example_models.StochasticDepthModel(drop_prob=0.5)
    model.train()
    x = torch.rand(2, 5)
    # layers_to_save triggers the two-pass path; use substring match
    mh = log_forward_pass(model, x, layers_to_save=["linear"])
    assert mh is not None
    assert len(mh.layer_labels) > 0
    # Linear layers should have saved activations
    linear_layers = [label for label in mh.layers_with_saved_activations if "linear" in label]
    assert len(linear_layers) > 0, "linear layers should have saved activations"


# =============================================================================
# Group A: Attention & Transformers
# =============================================================================


def test_multihead_attention(seq_input):
    model = example_models.MultiheadAttentionModel()
    assert validate_forward_pass(model, seq_input)
    show_model_graph(
        model,
        seq_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "multihead_attention"),
    )


def test_scaled_dot_product_attention():
    model = example_models.ScaledDotProductAttentionModel()
    x = torch.rand(2, 4, 10, 8)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "scaled_dot_product_attention"),
    )


def test_transformer_encoder(seq_input):
    model = example_models.TransformerEncoderModel()
    assert validate_forward_pass(model, seq_input)
    show_model_graph(
        model,
        seq_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "transformer_encoder"),
    )


def test_transformer_decoder():
    model = example_models.TransformerDecoderModel()
    tgt = torch.rand(5, 2, 16)
    memory = torch.rand(10, 2, 16)
    assert validate_forward_pass(model, (tgt, memory))
    show_model_graph(
        model,
        (tgt, memory),
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "transformer_decoder"),
    )


def test_embedding_positional(token_input):
    model = example_models.EmbeddingPositionalModel()
    assert validate_forward_pass(model, token_input)
    show_model_graph(
        model,
        token_input,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "embedding_positional"),
    )


def test_einsum():
    model = example_models.EinsumModel()
    x = torch.rand(2, 3, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "einsum"),
    )


# =============================================================================
# Group B: Container Modules
# =============================================================================


def test_module_list(input_2d):
    model = example_models.ModuleListModel()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_list"),
    )


def test_module_list_indexed(input_2d):
    model = example_models.ModuleListIndexedModel()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_list_indexed"),
    )


def test_module_dict(input_2d):
    model = example_models.ModuleDictModel()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "module_dict"),
    )


def test_var_args():
    model = example_models.VarArgsModel()
    inputs = [torch.rand(2, 5), torch.rand(2, 5), torch.rand(2, 5)]
    assert validate_forward_pass(model, inputs)
    show_model_graph(
        model,
        inputs,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "var_args"),
    )


def test_kwargs():
    model = example_models.KwargsModel()
    kwargs = {"a": torch.rand(2, 5), "b": torch.rand(2, 5)}
    assert validate_forward_pass(model, [], kwargs)
    show_model_graph(
        model,
        [],
        kwargs,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "kwargs"),
    )


# =============================================================================
# Group C: Conditional & Dynamic Ops
# =============================================================================


def test_torch_where():
    model = example_models.TorchWhereModel()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "torch_where"),
    )


def test_scatter_gather():
    model = example_models.ScatterGatherModel()
    x = torch.rand(3, 5)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "scatter_gather"),
    )


def test_no_grad_block():
    model = example_models.NoGradBlockModel()
    x = torch.rand(2, 32)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "no_grad_block"),
    )


def test_while_loop():
    model = example_models.WhileLoopModel()
    x = torch.rand(5, 5) * 10
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "while_loop"),
    )


def test_nested_conditional_loop():
    model = example_models.NestedConditionalLoopModel()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "nested_conditional_loop"),
    )


# =============================================================================
# Group D: Normalization & Conv Variants
# =============================================================================


def test_layer_norm():
    model = example_models.LayerNormModel()
    x = torch.rand(2, 3, 8)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "layer_norm"),
    )


def test_group_norm():
    model = example_models.GroupNormModel()
    x = torch.rand(2, 8, 4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "group_norm"),
    )


def test_instance_norm():
    model = example_models.InstanceNormModel()
    x = torch.rand(2, 3, 8, 8)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "instance_norm"),
    )


def test_conv1d(input_1d_seq):
    model = example_models.Conv1dModel()
    assert validate_forward_pass(model, input_1d_seq)
    show_model_graph(
        model,
        input_1d_seq,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "conv1d"),
    )


def test_conv3d(input_3d):
    model = example_models.Conv3dModel()
    assert validate_forward_pass(model, input_3d)
    show_model_graph(
        model,
        input_3d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "conv3d"),
    )


# =============================================================================
# Group E: Residual & Parameter Sharing
# =============================================================================


def test_residual_block():
    model = example_models.ResidualBlockModel()
    x = torch.rand(2, 16, 8, 8)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "residual_block"),
    )


def test_shared_param_branch(input_2d):
    model = example_models.SharedParamBranchModel()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "shared_param_branch"),
    )


def test_model_calling_model():
    model = example_models.ModelCallingModelModel()
    x = torch.rand(2, 32)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "model_calling_model"),
    )


def test_bidirectional_gru():
    model = example_models.BidirectionalGRUModel()
    x = torch.rand(5, 2, 8)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "bidirectional_gru"),
    )


# =============================================================================
# Group F: In-Place & Type Operations
# =============================================================================


def test_in_place_chain():
    model = example_models.InPlaceChainModel()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "in_place_chain"),
    )


def test_type_cast_chain():
    model = example_models.TypeCastChainModel()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "type_cast_chain"),
    )


def test_like_ops():
    model = example_models.LikeOpsModel()
    x = torch.rand(3, 3)
    assert validate_forward_pass(model, x, random_seed=42)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        random_seed=42,
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "like_ops"),
    )


def test_multi_tensor_return():
    model = example_models.MultiTensorReturnModel()
    x = torch.rand(3, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "multi_tensor_return"),
    )


def test_mixed_dtype():
    model = example_models.MixedDtypeModel()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "mixed_dtype"),
    )


# =============================================================================
# Group G: Scalar & Broadcasting
# =============================================================================


def test_scalar_tensor():
    model = example_models.ScalarTensorModel()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "scalar_tensor"),
    )


def test_broadcasting():
    model = example_models.BroadcastingModel()
    x = torch.rand(3, 4, 5)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "broadcasting"),
    )


def test_packed_sequence():
    model = example_models.PackedSequenceModel()
    x = torch.rand(5, 3, 8)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "packed_sequence"),
    )


def test_custom_autograd():
    model = example_models.CustomAutogradModel()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "custom_autograd"),
    )


# =============================================================================
# Group H: Exemption Registry Stress Tests
# =============================================================================


def test_cross_entropy():
    model = example_models.CrossEntropyModel()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "cross_entropy"),
    )


def test_index_select():
    model = example_models.IndexSelectModel()
    x = torch.rand(5, 3)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "index_select"),
    )


def test_interpolate():
    model = example_models.InterpolateModel()
    x = torch.rand(1, 1, 4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "interpolate"),
    )


def test_forward_hooks(input_2d):
    model = example_models.ForwardHooksModel()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "forward_hooks"),
    )


# =============================================================================
# Group I: Architecture Patterns
# =============================================================================


def test_simple_vae():
    model = example_models.SimpleVAE()
    x = torch.rand(2, 1, 28, 28)
    assert validate_forward_pass(model, x, random_seed=42)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        random_seed=42,
        vis_outpath=opj(VIS_OUTPUT_DIR, "generative-models", "simple_vae"),
    )


def test_simple_generator():
    model = example_models.SimpleGenerator()
    z = torch.randn(2, 100)
    assert validate_forward_pass(model, z)
    show_model_graph(
        model,
        z,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "generative-models", "simple_generator"),
    )


def test_simple_discriminator():
    model = example_models.SimpleDiscriminator()
    x = torch.rand(2, 1, 28, 28)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "generative-models", "simple_discriminator"),
    )


def test_small_unet():
    model = example_models.SmallUNet()
    x = torch.rand(1, 1, 32, 32)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "small_unet"),
    )


def test_temporal_conv_net():
    model = example_models.TemporalConvNet()
    x = torch.rand(2, 3, 32)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "time-series", "temporal_conv_net"),
    )


def test_espcn_super_res():
    model = example_models.ESPCNSuperRes()
    x = torch.rand(1, 1, 16, 16)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "super-resolution", "espcn_super_res"),
    )


def test_simple_pointnet():
    model = example_models.SimplePointNet()
    x = torch.rand(2, 3, 128)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "point-cloud", "simple_pointnet"),
    )


def test_actor_critic():
    model = example_models.ActorCritic()
    x = torch.rand(2, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "actor_critic"),
    )


def test_two_tower_recommender():
    model = example_models.TwoTowerRecommender()
    user = torch.rand(2, 8)
    item = torch.rand(2, 8)
    assert validate_forward_pass(model, (user, item))
    show_model_graph(
        model,
        (user, item),
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "two_tower_recommender"),
    )


def test_simple_depth_estimator():
    model = example_models.SimpleDepthEstimator()
    x = torch.rand(1, 3, 64, 64)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "simple_depth_estimator"),
    )


# =============================================================================
# Group Z: Adversarial Edge Cases
# =============================================================================


def test_same_tensor_all_args():
    model = example_models.SameTensorAllArgs()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "same_tensor_all_args"),
    )


def test_view_chain_mutate_middle():
    model = example_models.ViewChainMutateMiddle()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "view_chain_mutate_middle"),
    )


def test_self_caching_model(input_2d):
    model = example_models.SelfCachingModel()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "self_caching_model"),
    )


def test_delete_tensor_mid_forward(input_2d):
    model = example_models.DeleteTensorMidForward()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "delete_tensor_mid_forward"),
    )


def test_dynamic_module_creation(input_2d):
    model = example_models.DynamicModuleCreation()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "dynamic_module_creation"),
    )


def test_non_persistent_buffer(input_2d):
    model = example_models.NonPersistentBuffer()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "non_persistent_buffer"),
    )


def test_contiguous_no_op(input_2d):
    model = example_models.ContiguousNoOp()
    assert validate_forward_pass(model, input_2d)
    show_model_graph(
        model,
        input_2d,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "contiguous_no_op"),
    )


def test_stack_views_of_same_tensor():
    model = example_models.StackViewsOfSameTensor()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "stack_views_same_tensor"),
    )


def test_saturated_softmax():
    model = example_models.SaturatedSoftmax()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "saturated_softmax"),
    )


def test_duplicate_value_parents():
    model = example_models.DuplicateValueParents()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "duplicate_value_parents"),
    )


def test_empty_tensor_chain():
    model = example_models.EmptyTensorChain()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "empty_tensor_chain"),
    )


def test_clamp_narrow_range():
    model = example_models.ClampNarrowRange()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "clamp_narrow_range"),
    )


def test_round_near_integers():
    model = example_models.RoundNearIntegers()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "round_near_integers"),
    )


def test_bool_cast_exploit():
    model = example_models.BoolCastExploit()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "bool_cast_exploit"),
    )


def test_fake_loop_same_op_type():
    model = example_models.FakeLoopSameOpType()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "fake_loop_same_op_type"),
    )


def test_autocast_mid_forward():
    model = example_models.AutocastMidForward()
    x = torch.rand(4, 4)
    assert validate_forward_pass(model, x)
    show_model_graph(
        model,
        x,
        save_only=True,
        vis_opt="unrolled",
        vis_outpath=opj(VIS_OUTPUT_DIR, "toy-networks", "autocast_mid_forward"),
    )
