jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_getitem_with_transforms(self, array_dataset):
        assert len(array_dataset[0]) == 4
        assert len(array_dataset[1]) == 4
        assert len(array_dataset[2]) == 4
        assert len(array_dataset[3]) == 4
        torch.testing.assert_close(array_dataset[0][0], torch.tensor([1, 0, -1, 2]))
        torch.testing.assert_close(array_dataset[0][1], torch.tensor(1))
        np.testing.assert_array_equal(array_dataset[0][2], np.array([2, 1, -5, 1]))
        assert array_dataset[0][3] == 1
        torch.testing.assert_close(array_dataset[1][0], torch.tensor([1, 0, -2, -1]))
        torch.testing.assert_close(array_dataset[1][1], torch.tensor(0))
        np.testing.assert_array_equal(array_dataset[1][2], np.array([1, 0, -2, -1]))
        assert array_dataset[1][3] == 0
        torch.testing.assert_close(array_dataset[2][0], torch.tensor([2, 5, 0, 3]))
        torch.testing.assert_close(array_dataset[2][1], torch.tensor(0))
        np.testing.assert_array_equal(array_dataset[2][2], np.array([2, 5, 0, 3]))
        assert array_dataset[2][3] == 1
        torch.testing.assert_close(array_dataset[3][0], torch.tensor([-7, 1, 2, 2]))
        torch.testing.assert_close(array_dataset[3][1], torch.tensor(1))
        np.testing.assert_array_equal(array_dataset[3][2], np.array([-7, 1, 2, 2]))
        assert array_dataset[3][3] == 1
//  [python]
def test_process_temperature(inplace):
    from lmdeploy.pytorch.engine.logits_process import _process_temperature

    batch_size = 4
    num_tokens = 16
    scores = torch.rand(batch_size, num_tokens)
    temperatures = torch.rand(batch_size)

    gt = []
    for score, temperature in zip(scores, temperatures):
        warper = TemperatureLogitsWarper(temperature.item())
        gt.append(warper(None, score[None]))
    gt = torch.cat(gt)

    out = _process_temperature(scores, temperatures, inplace=inplace)
    torch.testing.assert_allclose(out, gt)
//  [python]
def test_process_temperature(inplace):
    from lmdeploy.pytorch.engine.logits_process import _process_temperature

    batch_size = 4
    num_tokens = 16
    scores = torch.rand(batch_size, num_tokens)
    temperatures = torch.rand(batch_size)

    gt = []
    for score, temperature in zip(scores, temperatures):
        warper = TemperatureLogitsWarper(temperature.item())
        gt.append(warper(None, score[None]))
    gt = torch.cat(gt)

    out = _process_temperature(scores, temperatures, inplace=inplace)
    torch.testing.assert_close(out, gt)
//  [python]
def test_pack_unpack_state(self):
        sm = TestScript.DerivedStateModule()
        x = torch.rand(3, 4)
        self.assertEqual(sm(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))

        # Test save path
        self.assertFalse(sm.pack_called.item())
        self.assertFalse(sm.unpack_called.item())
        imported = self.getExportImportCopyWithPacking(sm)
        # ensure pack was called before serialization
        self.assertTrue(sm.pack_called.item())
        # ensure unpack was called after serialization so as to leave the module in an initialized state
        self.assertTrue(sm.unpack_called.item())

        self.assertEqual(sm.derived, torch.neg(sm.param))

        # Test load paths
        self.assertTrue(imported.unpack_called.item())
        self.assertEqual(imported(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))
//  [python]
def test_pack_unpack_state(self):
        sm = TestScript.DerivedStateModule()
        x = torch.rand(3, 4)
        torch.testing.assert_close(sm(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))

        # Test save path
        self.assertFalse(sm.pack_called.item())
        self.assertFalse(sm.unpack_called.item())
        imported = self.getExportImportCopyWithPacking(sm)
        # ensure pack was called before serialization
        self.assertTrue(sm.pack_called.item())
        # ensure unpack was called after serialization so as to leave the module in an initialized state
        self.assertTrue(sm.unpack_called.item())

        torch.testing.assert_close(sm.derived, torch.neg(sm.param))

        # Test load paths
        self.assertTrue(imported.unpack_called.item())
        torch.testing.assert_close(imported(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))
//  [python]
def test_processrepetition_penalty(inplace):
    from lmdeploy.pytorch.engine.logits_process import \
        _process_repetition_penalty
    batch_size = 4
    num_tokens = 16
    scores = torch.rand(batch_size, num_tokens)
    input_ids = torch.tensor([
        [0, 1],
        [3, 6],
        [4, 4],
        [0, 0],
    ])
    penalties = 1 + torch.rand(batch_size)

    gt = []
    for score, ids, penalty in zip(scores, input_ids, penalties):
        warper = RepetitionPenaltyLogitsProcessor(penalty.item())
        gt.append(warper(ids[None], score[None].clone()))
    gt = torch.cat(gt)

    out = _process_repetition_penalty(scores,
                                      input_ids,
                                      penalties,
                                      inplace=inplace)
    t.assert_allclose(out, gt)
//  [python]
def test_processrepetition_penalty(inplace):
    from lmdeploy.pytorch.engine.logits_process import \
        _process_repetition_penalty
    batch_size = 4
    num_tokens = 16
    scores = torch.rand(batch_size, num_tokens)
    input_ids = torch.tensor([
        [0, 1],
        [3, 6],
        [4, 4],
        [0, 0],
    ])
    penalties = 1 + torch.rand(batch_size)

    gt = []
    for score, ids, penalty in zip(scores, input_ids, penalties):
        warper = RepetitionPenaltyLogitsProcessor(penalty.item())
        gt.append(warper(ids[None], score[None].clone()))
    gt = torch.cat(gt)

    out = _process_repetition_penalty(scores,
                                      input_ids,
                                      penalties,
                                      inplace=inplace)
    torch.testing.assert_close(out, gt)
//  [python]
def test_remove_dropout(self):
        weight_0_shape = (20, 5)
        weight_1_shape = (20, 20)
        input_shape = (10, 5)

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.weight_0 = torch.nn.Parameter(torch.rand(weight_0_shape))
                self.weight_1 = torch.nn.Parameter(torch.rand(weight_1_shape))

            def forward(self, x):
                o = F.linear(x, self.weight_0)
                o = F.dropout(o, training=self.training)
                o = F.linear(o, self.weight_1)
                return o

        data = torch.rand(input_shape)
        m = M()
        m = torch.jit.script(m)
        with self.assertRaisesRegex(RuntimeError, r'Dropout removal module in training mode is not yet supported'):
            torch._C._jit_pass_remove_dropout(m._c)
        m.eval()
        ref_res = m(data)
        # Need to inline otherwise we see instances of Function.
        # We would have to use torch.linear/dropout to get around it otherwise.
        from torch.jit._recursive import wrap_cpp_module
        m = wrap_cpp_module(torch._C._freeze_module(m._c))
        torch._C._jit_pass_remove_dropout(m._c)
        res = m(data)
        FileCheck().check_not("aten::dropout").run(str(m.graph))
        self.assertEqual(ref_res, res)
//  [python]
def test_remove_dropout(self):
        weight_0_shape = (20, 5)
        weight_1_shape = (20, 20)
        input_shape = (10, 5)

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.weight_0 = torch.nn.Parameter(torch.rand(weight_0_shape))
                self.weight_1 = torch.nn.Parameter(torch.rand(weight_1_shape))

            def forward(self, x):
                o = F.linear(x, self.weight_0)
                o = F.dropout(o, training=self.training)
                o = F.linear(o, self.weight_1)
                return o

        data = torch.rand(input_shape)
        m = M()
        m = torch.jit.script(m)
        with self.assertRaisesRegex(RuntimeError, r'Dropout removal module in training mode is not yet supported'):
            torch._C._jit_pass_remove_dropout(m._c)
        m.eval()
        ref_res = m(data)
        # Need to inline otherwise we see instances of Function.
        # We would have to use torch.linear/dropout to get around it otherwise.
        from torch.jit._recursive import wrap_cpp_module
        m = wrap_cpp_module(torch._C._freeze_module(m._c))
        torch._C._jit_pass_remove_dropout(m._c)
        res = m(data)
        FileCheck().check_not("aten::dropout").run(str(m.graph))
        torch.testing.assert_close(ref_res, res, rtol=1e-2, atol=1e-3)
//  [python]
def test_sparse_linearization():
    objective, ordering, A, b = build_test_objective_and_linear_system()

    linearization = th.SparseLinearization(objective, ordering=ordering)
    linearization.linearize()

    assert b.isclose(linearization.b).all()

    batch_size = A.shape[0]

    for i in range(batch_size):
        csrAi = linearization.structure().csr_straight(linearization.A_val[i, :])
        t.assert_close(A[i], torch.Tensor(csrAi.todense()))

    for i in range(batch_size):
        t.assert_close(b[i], linearization.b[i])

    # Test Atb result
    atb_expected = A.transpose(1, 2).bmm(b.unsqueeze(2))
    atb_out = linearization.Atb
    t.assert_close(atb_expected, atb_out)

    # Test Av() with a random v
    rng = torch.Generator()
    rng.manual_seed(1009)
    for _ in range(20):
        v = torch.randn(A.shape[0], A.shape[2], 1)
        av_expected = A.bmm(v).squeeze(2)
        av = linearization.Av(v.squeeze(2))
        t.assert_close(av_expected, av)

        v = v.squeeze(2)
        scaled_v_expected = (A.transpose(1, 2).bmm(A)).diagonal(dim1=1, dim2=2) * v
        scaled_v = linearization.diagonal_scaling(v)
        t.assert_close(scaled_v_expected, scaled_v)
//  [python]
def test_sparse_linearization():
    objective, ordering, A, b = build_test_objective_and_linear_system()

    linearization = th.SparseLinearization(objective, ordering=ordering)
    linearization.linearize()

    assert b.isclose(linearization.b).all()

    batch_size = A.shape[0]

    for i in range(batch_size):
        csrAi = linearization.structure().csr_straight(linearization.A_val[i, :])
        torch.testing.assert_close(A[i], torch.Tensor(csrAi.todense()))

    for i in range(batch_size):
        torch.testing.assert_close(b[i], linearization.b[i])

    # Test Atb result
    atb_expected = A.transpose(1, 2).bmm(b.unsqueeze(2))
    atb_out = linearization.Atb
    torch.testing.assert_close(atb_expected, atb_out)

    # Test Av() with a random v
    rng = torch.Generator()
    rng.manual_seed(1009)
    for _ in range(20):
        v = torch.randn(A.shape[0], A.shape[2], 1)
        av_expected = A.bmm(v).squeeze(2)
        av = linearization.Av(v.squeeze(2))
        torch.testing.assert_close(av_expected, av)

        v = v.squeeze(2)
        scaled_v_expected = (A.transpose(1, 2).bmm(A)).diagonal(dim1=1, dim2=2) * v
        scaled_v = linearization.diagonal_scaling(v)
        torch.testing.assert_close(scaled_v_expected, scaled_v)
//  [python]
def test_filter_topk_sorted(inplace):
    from lmdeploy.pytorch.engine.logits_process import _filter_topk_sorted

    batch_size = 4
    num_tokens = 16
    scores = torch.rand(batch_size, num_tokens).sort(1, descending=True)[0]
    top_k = torch.randint(4, num_tokens - 4, (batch_size, ))

    gt = []
    for score, k in zip(scores, top_k):
        warper = TopKLogitsWarper(k.item())
        gt.append(warper(None, score[None].clone()))
    gt = torch.cat(gt)

    out = _filter_topk_sorted(scores, top_k, inplace=inplace)
    torch.testing.assert_allclose(out, gt)
//  [python]
def test_filter_topk_sorted(inplace):
    from lmdeploy.pytorch.engine.logits_process import _filter_topk_sorted

    batch_size = 4
    num_tokens = 16
    scores = torch.rand(batch_size, num_tokens).sort(1, descending=True)[0]
    top_k = torch.randint(4, num_tokens - 4, (batch_size, ))

    gt = []
    for score, k in zip(scores, top_k):
        warper = TopKLogitsWarper(k.item())
        gt.append(warper(None, score[None].clone()))
    gt = torch.cat(gt)

    out = _filter_topk_sorted(scores, top_k, inplace=inplace)
    torch.testing.assert_close(out, gt)
//  [python]
def test_filter_topp_sorted(inplace):
    from lmdeploy.pytorch.engine.logits_process import _filter_topp_sorted

    batch_size = 4
    num_tokens = 16
    scores = torch.rand(batch_size, num_tokens).sort(1, descending=True)[0]
    top_p = torch.rand(batch_size)

    gt = []
    for score, p in zip(scores, top_p):
        warper = TopPLogitsWarper(p.item())
        gt.append(warper(None, score[None].clone()))
    gt = torch.cat(gt)

    out = _filter_topp_sorted(scores, top_p, inplace=inplace)
    torch.testing.assert_allclose(out, gt)
//  [python]
def test_filter_topp_sorted(inplace):
    from lmdeploy.pytorch.engine.logits_process import _filter_topp_sorted

    batch_size = 4
    num_tokens = 16
    scores = torch.rand(batch_size, num_tokens).sort(1, descending=True)[0]
    top_p = torch.rand(batch_size)

    gt = []
    for score, p in zip(scores, top_p):
        warper = TopPLogitsWarper(p.item())
        gt.append(warper(None, score[None].clone()))
    gt = torch.cat(gt)

    out = _filter_topp_sorted(scores, top_p, inplace=inplace)
    torch.testing.assert_close(out, gt)
//  [python]
def assert_close_optim_state(self, test: torch.Tensor, ref: torch.Tensor) -> None:
        tolerance = 1.0e-4 if test.dtype == torch.float else 1.0e-2

        torch.testing.assert_allclose(
            test.float().cpu(),
            ref.float().cpu(),
            atol=tolerance,
            rtol=tolerance,
        )
//  [python]
def assert_close_optim_state(self, test: torch.Tensor, ref: torch.Tensor) -> None:
        tolerance = 1.0e-4 if test.dtype == torch.float else 1.0e-2

        torch.testing.assert_close(
            test.float().cpu(),
            ref.float().cpu(),
            atol=tolerance,
            rtol=tolerance,
        )
//  [python]
def test_audio_melspec_consistency():
    dnsmos_th = DNSMOS_local(None, None, use_gpu=False, convert_to_torch=True)
    dnsmos_onnx = DNSMOS_local(None, None, use_gpu=False, convert_to_torch=False)
    x = torch.randn(8000)
    audio_melspec_th = dnsmos_th.audio_melspec(x)
    audio_melspec_onnx = dnsmos_onnx.audio_melspec(x.numpy())
    torch.testing.assert_allclose(audio_melspec_th, torch.from_numpy(audio_melspec_onnx))
//  [python]
def test_audio_melspec_consistency():
    dnsmos_th = DNSMOS_local(None, None, use_gpu=False, convert_to_torch=True)
    dnsmos_onnx = DNSMOS_local(None, None, use_gpu=False, convert_to_torch=False)
    x = torch.randn(8000)
    audio_melspec_th = dnsmos_th.audio_melspec(x)
    audio_melspec_onnx = dnsmos_onnx.audio_melspec(x.numpy())
    torch.testing.assert_close(audio_melspec_th, torch.from_numpy(audio_melspec_onnx))
//  [python]
def test_scale(self):
        boxes = torch.tensor([
            [0, 0, 1, 1],
            [0, 10, 10, 20],
        ])
        dtype = boxes.dtype

        boxlist = BoxList(boxes.clone())
        yscale, xscale = 1, 1
        boxlist.scale(yscale, xscale)
        self.assertEqual(boxlist.boxes, boxes)

        boxlist = BoxList(boxes.clone())
        yscale, xscale = 5, 3
        boxlist.scale(yscale, xscale)
        x_expected = (boxes[:, [0, 2]] * xscale).to(dtype=dtype)
        y_expected = (boxes[:, [1, 3]] * yscale).to(dtype=dtype)
        self.assertEqual(boxlist.boxes[:, [0, 2]], x_expected)
        self.assertEqual(boxlist.boxes[:, [1, 3]], y_expected)

        boxlist = BoxList(boxes.clone())
        yscale, xscale = 0.5, 0.3
        boxlist.scale(yscale, xscale)
        x_expected = (boxes[:, [0, 2]] * xscale).to(dtype=dtype)
        y_expected = (boxes[:, [1, 3]] * yscale).to(dtype=dtype)
        self.assertEqual(boxlist.boxes[:, [0, 2]], x_expected)
        self.assertEqual(boxlist.boxes[:, [1, 3]], y_expected)
//  [python]
def test_scale(self):
        boxes = torch.tensor([
            [0, 0, 1, 1],
            [0, 10, 10, 20],
        ])
        dtype = boxes.dtype

        boxlist = BoxList(boxes.clone())
        yscale, xscale = 1, 1
        boxlist.scale(yscale, xscale)
        torch.testing.assert_close(boxlist.boxes, boxes)

        boxlist = BoxList(boxes.clone())
        yscale, xscale = 5, 3
        boxlist.scale(yscale, xscale)
        x_expected = (boxes[:, [0, 2]] * xscale).to(dtype=dtype)
        y_expected = (boxes[:, [1, 3]] * yscale).to(dtype=dtype)
        torch.testing.assert_close(boxlist.boxes[:, [0, 2]], x_expected)
        torch.testing.assert_close(boxlist.boxes[:, [1, 3]], y_expected)

        boxlist = BoxList(boxes.clone())
        yscale, xscale = 0.5, 0.3
        boxlist.scale(yscale, xscale)
        x_expected = (boxes[:, [0, 2]] * xscale).to(dtype=dtype)
        y_expected = (boxes[:, [1, 3]] * yscale).to(dtype=dtype)
        torch.testing.assert_close(boxlist.boxes[:, [0, 2]], x_expected)
        torch.testing.assert_close(boxlist.boxes[:, [1, 3]], y_expected)
//  [python]
def test_get_polyfit_val_consistency(is_personalized_MOS):
    dnsmos_th = DNSMOS_local(None, None, use_gpu=False, convert_to_torch=True)
    dnsmos_onnx = DNSMOS_local(None, None, use_gpu=False, convert_to_torch=False)
    sig, bak, ovr = torch.rand(3) * 5
    mos_sig1, mos_bak1, mos_ovr1 = dnsmos_th.get_polyfit_val(
        sig, bak, ovr, is_personalized_MOS
    )
    mos_sig2, mos_bak2, mos_ovr2 = dnsmos_onnx.get_polyfit_val(
        sig.item(), bak.item(), ovr.item(), is_personalized_MOS
    )
    torch.testing.assert_allclose(mos_sig1, torch.as_tensor(mos_sig2).float())
    torch.testing.assert_allclose(mos_bak1, torch.as_tensor(mos_bak2).float())
    torch.testing.assert_allclose(mos_ovr1, torch.as_tensor(mos_ovr2).float())
//  [python]
def test_get_polyfit_val_consistency(is_personalized_MOS):
    dnsmos_th = DNSMOS_local(None, None, use_gpu=False, convert_to_torch=True)
    dnsmos_onnx = DNSMOS_local(None, None, use_gpu=False, convert_to_torch=False)
    sig, bak, ovr = torch.rand(3) * 5
    mos_sig1, mos_bak1, mos_ovr1 = dnsmos_th.get_polyfit_val(
        sig, bak, ovr, is_personalized_MOS
    )
    mos_sig2, mos_bak2, mos_ovr2 = dnsmos_onnx.get_polyfit_val(
        sig.item(), bak.item(), ovr.item(), is_personalized_MOS
    )
    torch.testing.assert_close(mos_sig1, torch.as_tensor(mos_sig2).float())
    torch.testing.assert_close(mos_bak1, torch.as_tensor(mos_bak2).float())
    torch.testing.assert_close(mos_ovr1, torch.as_tensor(mos_ovr2).float())
//  [python]
def test_transformer_layer_norm_init(self):
        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)

        src_vocab = trg_vocab = self.vocab

        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        def check_layer_norm(m: nn.Module):
            for _, child in m.named_children():
                if isinstance(child, nn.LayerNorm):
                    torch.testing.assert_allclose(
                        child.weight, torch.ones([self.hidden_size])
                    )
                    torch.testing.assert_allclose(
                        child.bias, torch.zeros([self.hidden_size])
                    )
                else:
                    check_layer_norm(child)

        check_layer_norm(model)
//  [python]
def test_transformer_layer_norm_init(self):
        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)

        src_vocab = trg_vocab = self.vocab

        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        def check_layer_norm(m: nn.Module):
            for _, child in m.named_children():
                if isinstance(child, nn.LayerNorm):
                    torch.testing.assert_close(
                        child.weight, torch.ones([self.hidden_size])
                    )
                    torch.testing.assert_close(
                        child.bias, torch.zeros([self.hidden_size])
                    )
                else:
                    check_layer_norm(child)

        check_layer_norm(model)
//  [python]
def test_stft_enc_dec_streaming(n_fft, win_length, hop_length, onesided):
    input_audio = torch.randn((1, 16000))
    ilens = torch.LongTensor([16000])

    encoder = STFTEncoder(
        n_fft=n_fft, win_length=win_length, hop_length=hop_length, onesided=onesided
    )
    decoder = STFTDecoder(
        n_fft=n_fft, win_length=win_length, hop_length=hop_length, onesided=onesided
    )
    frames, flens = encoder(input_audio, ilens)
    wav, ilens = decoder(frames, ilens)

    splited = encoder.streaming_frame(input_audio)

    sframes = [encoder.forward_streaming(s) for s in splited]
    swavs = [decoder.forward_streaming(s) for s in sframes]
    merged = decoder.streaming_merge(swavs, ilens)

    if not (is_torch_1_9_plus and encoder.use_builtin_complex):
        sframes = torch_complex.cat(sframes, dim=1)
    else:
        sframes = torch.cat(sframes, dim=1)

    torch.testing.assert_allclose(sframes.real, frames.real)
    torch.testing.assert_allclose(sframes.imag, frames.imag)
    torch.testing.assert_allclose(wav, input_audio)
    torch.testing.assert_allclose(wav, merged)
//  [python]
def test_stft_enc_dec_streaming(n_fft, win_length, hop_length, onesided):
    input_audio = torch.randn((1, 16000))
    ilens = torch.LongTensor([16000])

    encoder = STFTEncoder(
        n_fft=n_fft, win_length=win_length, hop_length=hop_length, onesided=onesided
    )
    decoder = STFTDecoder(
        n_fft=n_fft, win_length=win_length, hop_length=hop_length, onesided=onesided
    )
    frames, flens = encoder(input_audio, ilens)
    wav, ilens = decoder(frames, ilens)

    splited = encoder.streaming_frame(input_audio)

    sframes = [encoder.forward_streaming(s) for s in splited]
    swavs = [decoder.forward_streaming(s) for s in sframes]
    merged = decoder.streaming_merge(swavs, ilens)

    if not (is_torch_1_9_plus and encoder.use_builtin_complex):
        sframes = torch_complex.cat(sframes, dim=1)
    else:
        sframes = torch.cat(sframes, dim=1)

    torch.testing.assert_close(sframes.real, frames.real)
    torch.testing.assert_close(sframes.imag, frames.imag)
    torch.testing.assert_close(wav, input_audio)
    torch.testing.assert_close(wav, merged)
//  [python]
def test_setattr(self):
        # some modules (LSTM) rewrite __setattr__ which may break the logic
        # these are tested here
        x = torch.randn(2, 3, 10)
        lstm = nn.LSTM(10, 11)
        y0, (h0, c0) = lstm(x)
        params = make_functional(lstm)
        y1, (h1, c1) = lstm(x, params=params)
        y2, (h2, c2) = lstm(x, (h1, c1), params=params)
        torch.testing.assert_allclose(y1, y0)
        params.apply_(lambda p: p.data.zero_())
        y1, (h1, c1) = lstm(x, params=params)
        assert not torch.isclose(y0, y1).all()
//  [python]
def test_setattr(self):
        # some modules (LSTM) rewrite __setattr__ which may break the logic
        # these are tested here
        x = torch.randn(2, 3, 10)
        lstm = nn.LSTM(10, 11)
        y0, (h0, c0) = lstm(x)
        params = make_functional(lstm)
        y1, (h1, c1) = lstm(x, params=params)
        y2, (h2, c2) = lstm(x, (h1, c1), params=params)
        torch.testing.assert_close(y1, y0)
        params.apply_(lambda p: p.data.zero_())
        y1, (h1, c1) = lstm(x, params=params)
        assert not torch.isclose(y0, y1).all()
//  [python]
def _test_gemm_rcr_bias_softmax(
        self,
        M=16,
        K=64,
        N=24,
        dtype="float16",
        test_name="gemm_rcr_bias_softmax",
        atol=1e-2,
        rtol=1e-2,
        assert_argmax=True,
    ):
        X = Tensor(shape=[M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[N, K], dtype=dtype, name="input_1", is_input=True)
        B = Tensor(shape=[N], dtype=dtype, name="input_2", is_input=True)
        OP = ops.gemm_rcr_bias_softmax()
        Y = OP(X, W, B)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        x_pt = get_random_torch_tensor([M, K], dtype)
        w_pt = get_random_torch_tensor([N, K], dtype)
        b_pt = get_random_torch_tensor([N], dtype)
        y_pt = torch.nn.functional.linear(x_pt, w_pt, bias=b_pt)
        y_pt = torch.softmax(y_pt, dim=1)

        module = compile_model(Y, detect_target(), "./tmp", test_name)

        inputs = {"input_0": x_pt, "input_1": w_pt, "input_2": b_pt}
        y = get_torch_empty_tensor([M, N], dtype)
        module.run_with_tensors(inputs, [y])

        torch.testing.assert_allclose(y, y_pt, atol=atol, rtol=rtol)

        if assert_argmax:
            torch.testing.assert_allclose(
                torch.argmax(y, axis=1),
                torch.argmax(y_pt, axis=1),
                atol=1e-1,
                rtol=1e-1,
            )
//  [python]
def _test_gemm_rcr_bias_softmax(
        self,
        M=16,
        K=64,
        N=24,
        dtype="float16",
        test_name="gemm_rcr_bias_softmax",
        atol=1e-2,
        rtol=1e-2,
        assert_argmax=True,
    ):
        X = Tensor(shape=[M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[N, K], dtype=dtype, name="input_1", is_input=True)
        B = Tensor(shape=[N], dtype=dtype, name="input_2", is_input=True)
        OP = ops.gemm_rcr_bias_softmax()
        Y = OP(X, W, B)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        x_pt = get_random_torch_tensor([M, K], dtype)
        w_pt = get_random_torch_tensor([N, K], dtype)
        b_pt = get_random_torch_tensor([N], dtype)
        y_pt = torch.nn.functional.linear(x_pt, w_pt, bias=b_pt)
        y_pt = torch.softmax(y_pt, dim=1)

        module = compile_model(Y, detect_target(), "./tmp", test_name)

        inputs = {"input_0": x_pt, "input_1": w_pt, "input_2": b_pt}
        y = get_torch_empty_tensor([M, N], dtype)
        module.run_with_tensors(inputs, [y])

        torch.testing.assert_close(y, y_pt, atol=atol, rtol=rtol)

        if assert_argmax:
            torch.testing.assert_close(
                torch.argmax(y, axis=1),
                torch.argmax(y_pt, axis=1),
                atol=1e-1,
                rtol=1e-1,
            )
//  [python]
def recursive_value_check(out, ref_out):
    if isinstance(out, torch.Tensor):
        assert isinstance(ref_out, torch.Tensor)
        torch.allclose(out, ref_out)
    elif isinstance(out, (tuple, list)):
        assert isinstance(ref_out, type(out))
        for a, b in zip(out, ref_out):
            recursive_value_check(a, b)
    elif isinstance(out, dict):
        assert isinstance(ref_out, dict)
        assert set(out.keys()) == set(ref_out.keys())
        for k in out.keys():
            recursive_value_check(out[k], ref_out[k])
    else:
        raise RuntimeError(f"Unsupported type {type(out)}")
//  [python]
def recursive_value_check(out, ref_out):
    if isinstance(out, torch.Tensor):
        assert isinstance(ref_out, torch.Tensor)
        torch.testing.assert_close(out, ref_out)
    elif isinstance(out, (tuple, list)):
        assert isinstance(ref_out, type(out))
        for a, b in zip(out, ref_out):
            recursive_value_check(a, b)
    elif isinstance(out, dict):
        assert isinstance(ref_out, dict)
        assert set(out.keys()) == set(ref_out.keys())
        for k in out.keys():
            recursive_value_check(out[k], ref_out[k])
    else:
        raise RuntimeError(f"Unsupported type {type(out)}")
//  [python]
def test_transformer_layers_init(self):
        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)
        cfg["model"]["initializer"] = "xavier_normal"
        cfg["model"]["encoder"]["num_layers"] = 6
        cfg["model"]["decoder"]["num_layers"] = 6

        src_vocab = trg_vocab = self.vocab

        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        self.assertEqual(len(model.encoder.layers), 6)
        self.assertEqual(len(model.decoder.layers), 6)

        for layer in model.encoder.layers:
            self.assertEqual(layer.alpha, 1.417938140685523)
        for layer in model.decoder.layers:
            self.assertEqual(layer.alpha, 2.0597671439071177)

        self.assertTrue(torch.allclose(
            model.encoder.layers[0].src_src_att.q_layer.weight[:5, 0].data,
            torch.Tensor([0.1232, 0.1870, -0.1077, 0.0748, -0.0651]),
            rtol=1e-4,
            atol=1e-4,
        ))
        self.assertTrue(torch.allclose(
            model.decoder.layers[0].src_trg_att.q_layer.weight[:5, 0].data,
            torch.Tensor([-0.1035, 0.2171, 0.1729, -0.0120, -0.1008]),
            rtol=1e-4,
            atol=1e-4,
        ))
        self.assertTrue(torch.allclose(
            model.decoder.layers[0].trg_trg_att.q_layer.weight[:5, 0].data,
            torch.Tensor([-0.2248, -0.0396, 0.2041, 0.0627, 0.0255]),
            rtol=1e-4,
            atol=1e-4,
        ))
//  [python]
def test_transformer_layers_init(self):
        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)
        cfg["model"]["initializer"] = "xavier_normal"
        cfg["model"]["encoder"]["num_layers"] = 6
        cfg["model"]["decoder"]["num_layers"] = 6

        src_vocab = trg_vocab = self.vocab

        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        self.assertEqual(len(model.encoder.layers), 6)
        self.assertEqual(len(model.decoder.layers), 6)

        for layer in model.encoder.layers:
            self.assertEqual(layer.alpha, 1.417938140685523)
        for layer in model.decoder.layers:
            self.assertEqual(layer.alpha, 2.0597671439071177)

        torch.testing.assert_close(
            model.encoder.layers[0].src_src_att.q_layer.weight[:5, 0].data,
            torch.Tensor([0.1232, 0.1870, -0.1077, 0.0748, -0.0651]),
            rtol=1e-4,
            atol=1e-4,
        )
        torch.testing.assert_close(
            model.decoder.layers[0].src_trg_att.q_layer.weight[:5, 0].data,
            torch.Tensor([-0.1035, 0.2171, 0.1729, -0.0120, -0.1008]),
            rtol=1e-4,
            atol=1e-4,
        )
        torch.testing.assert_close(
            model.decoder.layers[0].trg_trg_att.q_layer.weight[:5, 0].data,
            torch.Tensor([-0.2248, -0.0396, 0.2041, 0.0627, 0.0255]),
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def _test_bmm_rcr_softmax(
        self,
        B=16,
        M=16,
        K=64,
        N=24,
        dtype="float16",
        test_name="bmm_rcr_softmax",
    ):
        X = Tensor(shape=[B, M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[B, N, K], dtype=dtype, name="input_1", is_input=True)
        OP = ops.bmm_rcr_softmax()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        x_pt = get_random_torch_tensor([B, M, K], dtype)
        w_pt = get_random_torch_tensor([B, N, K], dtype)
        wt_pt = torch.transpose(w_pt, 2, 1)
        y_pt = torch.bmm(x_pt, wt_pt)
        y_pt = torch.softmax(y_pt, dim=-1)

        module = compile_model(Y, detect_target(), "./tmp", test_name)

        inputs = {"input_0": x_pt, "input_1": w_pt}
        y = get_torch_empty_tensor([B, M, N], dtype)
        module.run_with_tensors(inputs, [y])

        torch.allclose(y, y_pt, atol=1e-2, rtol=1e-2)

        torch.allclose(
            torch.argmax(y, axis=2),
            torch.argmax(y_pt, axis=2),
            atol=1e-1,
            rtol=1e-1,
        )
//  [python]
def _test_bmm_rcr_softmax(
        self,
        B=16,
        M=16,
        K=64,
        N=24,
        dtype="float16",
        test_name="bmm_rcr_softmax",
    ):
        X = Tensor(shape=[B, M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[B, N, K], dtype=dtype, name="input_1", is_input=True)
        OP = ops.bmm_rcr_softmax()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        x_pt = get_random_torch_tensor([B, M, K], dtype)
        w_pt = get_random_torch_tensor([B, N, K], dtype)
        wt_pt = torch.transpose(w_pt, 2, 1)
        y_pt = torch.bmm(x_pt, wt_pt)
        y_pt = torch.softmax(y_pt, dim=-1)

        module = compile_model(Y, detect_target(), "./tmp", test_name)

        inputs = {"input_0": x_pt, "input_1": w_pt}
        y = get_torch_empty_tensor([B, M, N], dtype)
        module.run_with_tensors(inputs, [y])

        torch.testing.assert_close(y, y_pt, atol=1e-2, rtol=1e-2)

        torch.testing.assert_close(
            torch.argmax(y, axis=2),
            torch.argmax(y_pt, axis=2),
            atol=1e-1,
            rtol=1e-1,
        )
//  [python]
def _check_fx_compatible(model, inputs, eager_out=None):
    model_fx = torch.fx.symbolic_trace(model)
    if eager_out is None:
        eager_out = model(inputs)
    with torch.no_grad(), freeze_rng_state():
        fx_out = model_fx(inputs)
    torch.testing.assert_allclose(eager_out, fx_out)
//  [python]
def _check_fx_compatible(model, inputs, eager_out=None):
    model_fx = torch.fx.symbolic_trace(model)
    if eager_out is None:
        eager_out = model(inputs)
    with torch.no_grad(), freeze_rng_state():
        fx_out = model_fx(inputs)
    torch.testing.assert_close(eager_out, fx_out)
//  [python]
def _assert_expected(output, name, prec=None, atol=None, rtol=None):
    """Test that a python value matches the recorded contents of a file
    based on a "check" name. The value must be
    pickable with `torch.save`. This file
    is placed in the 'expect' directory in the same directory
    as the test script. You can automatically update the recorded test
    output using an EXPECTTEST_ACCEPT=1 env variable.
    """
    expected_file = _get_expected_file(name)

    if ACCEPT:
        filename = {os.path.basename(expected_file)}
        print(f"Accepting updated output for {filename}:\n\n{output}")
        torch.save(output, expected_file)
        MAX_PICKLE_SIZE = 50 * 1000  # 50 KB
        binary_size = os.path.getsize(expected_file)
        if binary_size > MAX_PICKLE_SIZE:
            raise RuntimeError(f"The output for {filename}, is larger than 50kb - got {binary_size}kb")
    else:
        expected = torch.load(expected_file, weights_only=True)
        rtol = rtol or prec  # keeping prec param for legacy reason, but could be removed ideally
        atol = atol or prec
        t.testing.assert_allclose(output, expected, rtol=rtol, atol=atol, check_dtype=False, check_device=False)
//  [python]
def _assert_expected(output, name, prec=None, atol=None, rtol=None):
    """Test that a python value matches the recorded contents of a file
    based on a "check" name. The value must be
    pickable with `torch.save`. This file
    is placed in the 'expect' directory in the same directory
    as the test script. You can automatically update the recorded test
    output using an EXPECTTEST_ACCEPT=1 env variable.
    """
    expected_file = _get_expected_file(name)

    if ACCEPT:
        filename = {os.path.basename(expected_file)}
        print(f"Accepting updated output for {filename}:\n\n{output}")
        torch.save(output, expected_file)
        MAX_PICKLE_SIZE = 50 * 1000  # 50 KB
        binary_size = os.path.getsize(expected_file)
        if binary_size > MAX_PICKLE_SIZE:
            raise RuntimeError(f"The output for {filename}, is larger than 50kb - got {binary_size}kb")
    else:
        expected = torch.load(expected_file, weights_only=True)
        rtol = rtol or prec  # keeping prec param for legacy reason, but could be removed ideally
        atol = atol or prec
        torch.testing.assert_close(output, expected, rtol=rtol, atol=atol, check_dtype=False, check_device=False)
//  [python]
def _check_jit_scriptable(nn_module, args, unwrapper=None, eager_out=None):
    """Check that a nn.Module's results in TorchScript match eager and that it can be exported"""

    def get_export_import_copy(m):
        """Save and load a TorchScript model"""
        with TemporaryDirectory() as dir:
            path = os.path.join(dir, "script.pt")
            m.save(path)
            imported = torch.jit.load(path)
        return imported

    sm = torch.jit.script(nn_module)
    sm.eval()

    if eager_out is None:
        with torch.no_grad(), freeze_rng_state():
            eager_out = nn_module(*args)

    with torch.no_grad(), freeze_rng_state():
        script_out = sm(*args)
        if unwrapper:
            script_out = unwrapper(script_out)

    t.assert_close(eager_out, script_out, atol=1e-4, rtol=1e-4)

    m_import = get_export_import_copy(sm)
    with torch.no_grad(), freeze_rng_state():
        imported_script_out = m_import(*args)
        if unwrapper:
            imported_script_out = unwrapper(imported_script_out)

    t.assert_close(script_out, imported_script_out, atol=3e-4, rtol=3e-4)
//  [python]
def _check_jit_scriptable(nn_module, args, unwrapper=None, eager_out=None):
    """Check that a nn.Module's results in TorchScript match eager and that it can be exported"""

    def get_export_import_copy(m):
        """Save and load a TorchScript model"""
        with TemporaryDirectory() as dir:
            path = os.path.join(dir, "script.pt")
            m.save(path)
            imported = torch.jit.load(path)
        return imported

    sm = torch.jit.script(nn_module)
    sm.eval()

    if eager_out is None:
        with torch.no_grad(), freeze_rng_state():
            eager_out = nn_module(*args)

    with torch.no_grad(), freeze_rng_state():
        script_out = sm(*args)
        if unwrapper:
            script_out = unwrapper(script_out)

    torch.testing.assert_close(eager_out, script_out, atol=1e-4, rtol=1e-4)

    m_import = get_export_import_copy(sm)
    with torch.no_grad(), freeze_rng_state():
        imported_script_out = m_import(*args)
        if unwrapper:
            imported_script_out = unwrapper(imported_script_out)

    torch.testing.assert_close(script_out, imported_script_out, atol=3e-4, rtol=3e-4)
//  [python]
def test_memory_efficient_densenet(model_fn):
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

    model1 = model_fn(num_classes=50, memory_efficient=True)
    params = model1.state_dict()
    num_params = sum(x.numel() for x in model1.parameters())
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
    num_grad = sum(x.grad.numel() for x in model1.parameters() if x.grad is not None)

    model2 = model_fn(num_classes=50, memory_efficient=False)
    model2.load_state_dict(params)
    model2.eval()
    out2 = model2(x)

    assert num_params == num_grad
    torch.testing.assert_allclose(out1, out2, rtol=0.0, atol=1e-5)

    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)
//  [python]
def test_memory_efficient_densenet(model_fn):
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

    model1 = model_fn(num_classes=50, memory_efficient=True)
    params = model1.state_dict()
    num_params = sum(x.numel() for x in model1.parameters())
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
    num_grad = sum(x.grad.numel() for x in model1.parameters() if x.grad is not None)

    model2 = model_fn(num_classes=50, memory_efficient=False)
    model2.load_state_dict(params)
    model2.eval()
    out2 = model2(x)

    assert num_params == num_grad
    torch.testing.assert_close(out1, out2, rtol=0.0, atol=1e-5)

    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)
//  [python]
def assert_close(actual, expected, *args, **kwargs):
    if hasattr(torch.testing, 'assert_close'):
        torch.testing.assert_close(actual, expected, *args, **kwargs)
    else:
        filtered_kwargs = {k: v for k, v in kwargs.items() if not k.startswith('check_')}
        if 'msg' in filtered_kwargs:
            msg = filtered_kwargs['msg']
            if not isinstance(msg, str):
                filtered_kwargs['msg'] = msg()

        if not isinstance(actual, (tuple, list)):
            actual_v = [actual]
            expected_v = [expected]
        else:
            actual_v = actual
            expected_v = expected

        for a, e in zip(actual_v, expected_v):
            if not kwargs.get('check_dtype', True):
                t = torch.promote_types(a.dtype, e.dtype)
                a = a.to(dtype=t)
                e = e.to(dtype=t)
            torch.testing.assert_allclose(a, e, *args, **filtered_kwargs)
//  [python]
def test_parallelism_no_zero_3(devices, model_name):
    model = AutoModel.from_pretrained(model_name).to(devices[0])
    model_state_dict = model.state_dict()
    model_tp = TensorParallelPreTrainedModel(model, devices, sharded=False)
    del model
    with save_tensor_parallel(model_tp):
        model_tp_state_dict = model_tp.state_dict()
    del model_tp

    assert sorted(list(model_state_dict.keys())) == sorted(list(model_tp_state_dict.keys()))

    for name in model_state_dict.keys():
        data = model_state_dict[name]
        data_tp = model_tp_state_dict[name]

        assert data.shape == data_tp.shape

        torch.testing.assert_allclose(data, data_tp)
//  [python]
def test_parallelism_no_zero_3(devices, model_name):
    model = AutoModel.from_pretrained(model_name).to(devices[0])
    model_state_dict = model.state_dict()
    model_tp = TensorParallelPreTrainedModel(model, devices, sharded=False)
    del model
    with save_tensor_parallel(model_tp):
        model_tp_state_dict = model_tp.state_dict()
    del model_tp

    assert sorted(list(model_state_dict.keys())) == sorted(list(model_tp_state_dict.keys()))

    for name in model_state_dict.keys():
        data = model_state_dict[name]
        data_tp = model_tp_state_dict[name]

        assert data.shape == data_tp.shape

        torch.testing.assert_close(data, data_tp)
//  [python]
def test_no_parallelism_zero_3(devices, model_name):
    model = AutoModel.from_pretrained(model_name).to(devices[0])
    model_state_dict = model.state_dict()
    model_tp = TensorParallel(
        model, devices, tensor_parallel_config=Config({}, {}, {}, {}), sharded=True
    )  # zero-3 sharding only
    del model
    with save_tensor_parallel(model_tp):
        model_tp_state_dict = model_tp.state_dict()
    del model_tp

    assert sorted(list(model_state_dict.keys())) == sorted(list(model_tp_state_dict.keys()))

    for name in model_state_dict.keys():
        data = model_state_dict[name]
        data_tp = model_tp_state_dict[name]

        assert data.shape == data_tp.shape

        torch.testing.assert_allclose(data, data_tp)
//  [python]
def test_no_parallelism_zero_3(devices, model_name):
    model = AutoModel.from_pretrained(model_name).to(devices[0])
    model_state_dict = model.state_dict()
    model_tp = TensorParallel(
        model, devices, tensor_parallel_config=Config({}, {}, {}, {}), sharded=True
    )  # zero-3 sharding only
    del model
    with save_tensor_parallel(model_tp):
        model_tp_state_dict = model_tp.state_dict()
    del model_tp

    assert sorted(list(model_state_dict.keys())) == sorted(list(model_tp_state_dict.keys()))

    for name in model_state_dict.keys():
        data = model_state_dict[name]
        data_tp = model_tp_state_dict[name]

        assert data.shape == data_tp.shape

        torch.testing.assert_close(data, data_tp)
//  [python]
def test_parallelism_zero_3(devices, model_name):
    model = AutoModel.from_pretrained(model_name).to(devices[0])
    model_state_dict = model.state_dict()
    model_tp = TensorParallelPreTrainedModel(model, devices, sharded=True)
    del model
    with save_tensor_parallel(model_tp):
        model_tp_state_dict = model_tp.state_dict()
    del model_tp

    assert sorted(list(model_state_dict.keys())) == sorted(list(model_tp_state_dict.keys()))

    for name in model_state_dict.keys():
        data = model_state_dict[name]
        data_tp = model_tp_state_dict[name]

        assert data.shape == data_tp.shape, name

        torch.testing.assert_allclose(data, data_tp, msg=lambda msg: f"{name}:\n{msg}")
//  [python]
def test_parallelism_zero_3(devices, model_name):
    model = AutoModel.from_pretrained(model_name).to(devices[0])
    model_state_dict = model.state_dict()
    model_tp = TensorParallelPreTrainedModel(model, devices, sharded=True)
    del model
    with save_tensor_parallel(model_tp):
        model_tp_state_dict = model_tp.state_dict()
    del model_tp

    assert sorted(list(model_state_dict.keys())) == sorted(list(model_tp_state_dict.keys()))

    for name in model_state_dict.keys():
        data = model_state_dict[name]
        data_tp = model_tp_state_dict[name]

        assert data.shape == data_tp.shape, name

        torch.testing.assert_close(data, data_tp, msg=lambda msg: f"{name}:\n{msg}")
//  [python]
torch.save(split_mod, "trt.pt")
reload_trt_mod = torch.load("trt.pt")
reload_model_output = reload_trt_mod(*inputs)

# Make sure the results match
regular_model_output = model(*inputs)
torch.allclose(
    reload_model_output, regular_model_output, atol=3e-3, rtol=1e-2
)
//  [python]
torch.save(split_mod, "trt.pt")
reload_trt_mod = torch.load("trt.pt")
reload_model_output = reload_trt_mod(*inputs)

# Make sure the results match
regular_model_output = model(*inputs)
torch.testing.assert_close(
    reload_model_output, regular_model_output, atol=3e-3, rtol=1e-2
)
//  [python]
def test_selective_torch_compile(self):
        m = nn.Sequential(
            nn.Linear(4, 4),
            nn.Sequential(
                nn.Linear(4, 4),
                nn.Linear(4, 4),
            ),
            nn.Linear(4, 4),
        )
        x = torch.randn(4, 4)
        y_ref = m(x)

        _replace_with_custom_fn_if_matches_filter(
            m,
            lambda mod: torch.compile(mod),
            lambda mod, fqn: isinstance(mod, nn.Linear) and fqn != "1.0",
        )

        self.assertTrue(isinstance(m[0], torch._dynamo.eval_frame.OptimizedModule))
        self.assertTrue(isinstance(m[1][0], nn.Linear))
        self.assertTrue(isinstance(m[1][1], torch._dynamo.eval_frame.OptimizedModule))
        self.assertTrue(isinstance(m[2], torch._dynamo.eval_frame.OptimizedModule))

        y = m(x)
        torch.testing.assert_allclose(y, y_ref)
//  [python]
def test_selective_torch_compile(self):
        m = nn.Sequential(
            nn.Linear(4, 4),
            nn.Sequential(
                nn.Linear(4, 4),
                nn.Linear(4, 4),
            ),
            nn.Linear(4, 4),
        )
        x = torch.randn(4, 4)
        y_ref = m(x)

        _replace_with_custom_fn_if_matches_filter(
            m,
            lambda mod: torch.compile(mod),
            lambda mod, fqn: isinstance(mod, nn.Linear) and fqn != "1.0",
        )

        self.assertTrue(isinstance(m[0], torch._dynamo.eval_frame.OptimizedModule))
        self.assertTrue(isinstance(m[1][0], nn.Linear))
        self.assertTrue(isinstance(m[1][1], torch._dynamo.eval_frame.OptimizedModule))
        self.assertTrue(isinstance(m[2], torch._dynamo.eval_frame.OptimizedModule))

        y = m(x)
        torch.testing.assert_close(y, y_ref)