jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_save_and_load_quantized(self, backend: Backend) -> None:
        model = SimpleModule()
        quantizer = Quantizer(backend=backend)
        dummy_input = torch.randn((1, 2, 5, 5))
        calib_model = quantizer.calib(model)
        calib_model(dummy_input)

        quant_model = quantizer.quantize(model)
        quant_output = quant_model(dummy_input)
        ts_quant_model = torch.jit.trace(quant_model, dummy_input)

        with tempfile.NamedTemporaryFile() as tmp_file:
            torch.jit.save(ts_quant_model, tmp_file.name)
            loaded = torch.jit.load(tmp_file.name)
            loaded = torch.jit.load(tmp_file.name)
        loaded_output = loaded(dummy_input)
        torch.testing.assert_allclose(quant_output, loaded_output)
//  [python]
def test_save_and_load_quantized(self, backend: Backend) -> None:
        model = SimpleModule()
        quantizer = Quantizer(backend=backend)
        dummy_input = torch.randn((1, 2, 5, 5))
        calib_model = quantizer.calib(model)
        calib_model(dummy_input)

        quant_model = quantizer.quantize(model)
        quant_output = quant_model(dummy_input)
        ts_quant_model = torch.jit.trace(quant_model, dummy_input)

        with tempfile.NamedTemporaryFile() as tmp_file:
            torch.jit.save(ts_quant_model, tmp_file.name)
            loaded = torch.jit.load(tmp_file.name)
            loaded = torch.jit.load(tmp_file.name)
        loaded_output = loaded(dummy_input)
        torch.testing.assert_close(quant_output, loaded_output)
//  [python]
def test_scrubbing(model_str: str):
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

    # It turns out to be important to use the actual pretrained weights and not random
    # weights, since random models tend to get near-constant loss even if you implement
    # them incorrectly.
    model = AutoModelForCausalLM.from_pretrained(model_str).to(device)
    tokenizer = AutoTokenizer.from_pretrained(model_str)

    ## Test attention patching by itself ##
    # Get the output pre-patching
    input_ids = tokenizer.encode("Hello world", return_tensors="pt").to(device)
    clean_logits = model(input_ids).logits

    # Patch the model
    patch_attention_(model)

    # Check that the outputs are the same
    patched_logits = model(input_ids).logits
    torch.testing.assert_allclose(clean_logits, patched_logits, atol=5e-5, rtol=5e-5)

    ## Test attention patching + concept erasure ##
    tokenized = tokenizer(
        ["The quick brown fox jumps over the lazy dog", "Hello world!"]
    )
    ds = Dataset.from_dict(dict(tokenized)).with_format("torch")
    gt_losses = []

    for record in ds:
        input_ids = record["input_ids"].unsqueeze(0).to(device)
        gt_losses.append(model(input_ids, labels=input_ids).loss)

    gt_loss = torch.stack(gt_losses).mean().item()
    _, loss = scrub(model, ds, z_column=None)
    assert math.isclose(gt_loss, loss, rel_tol=1e-4)
//  [python]
def test_scrubbing(model_str: str):
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

    # It turns out to be important to use the actual pretrained weights and not random
    # weights, since random models tend to get near-constant loss even if you implement
    # them incorrectly.
    model = AutoModelForCausalLM.from_pretrained(model_str).to(device)
    tokenizer = AutoTokenizer.from_pretrained(model_str)

    ## Test attention patching by itself ##
    # Get the output pre-patching
    input_ids = tokenizer.encode("Hello world", return_tensors="pt").to(device)
    clean_logits = model(input_ids).logits

    # Patch the model
    patch_attention_(model)

    # Check that the outputs are the same
    patched_logits = model(input_ids).logits
    torch.testing.assert_close(clean_logits, patched_logits, atol=5e-5, rtol=5e-5)

    ## Test attention patching + concept erasure ##
    tokenized = tokenizer(
        ["The quick brown fox jumps over the lazy dog", "Hello world!"]
    )
    ds = Dataset.from_dict(dict(tokenized)).with_format("torch")
    gt_losses = []

    for record in ds:
        input_ids = record["input_ids"].unsqueeze(0).to(device)
        gt_losses.append(model(input_ids, labels=input_ids).loss)

    gt_loss = torch.stack(gt_losses).mean().item()
    _, loss = scrub(model, ds, z_column=None)
    assert math.isclose(gt_loss, loss, rel_tol=1e-4)
//  [python]
def test_calib_and_quantize_with_op_types_filter(
        self, backend: Backend, module_filter: ModuleFilter
    ) -> None:
        model = SimpleModule()
        quantizer = Quantizer(backend=backend, module_filter=module_filter)
        dummy_input = torch.randn((1, 2, 5, 5))
        original_output = model(dummy_input)

        calib_model = quantizer.calib(model)
        calib_output = calib_model(dummy_input)
        self.assertTrue(torch.equal(original_output, calib_output))

        quant_model = quantizer.quantize(model)
        self.assertTrue(isinstance(quant_model.conv, torch.nn.Conv2d))
        self.assertTrue(isinstance(quant_model.sub.conv, torch.nn.Conv2d))
        self.assertEqual(quant_model.conv, model.conv)
        self.assertEqual(quant_model.sub.conv, model.sub.conv)
        self.assertNotEqual(quant_model.linear, model.linear)

        quant_output = quant_model(dummy_input)
        self.assertFalse(torch.equal(original_output, quant_output))
        torch.allclose(quant_output, original_output, rtol=0.1, atol=0.5)
//  [python]
def test_calib_and_quantize_with_op_types_filter(
        self, backend: Backend, module_filter: ModuleFilter
    ) -> None:
        model = SimpleModule()
        quantizer = Quantizer(backend=backend, module_filter=module_filter)
        dummy_input = torch.randn((1, 2, 5, 5))
        original_output = model(dummy_input)

        calib_model = quantizer.calib(model)
        calib_output = calib_model(dummy_input)
        self.assertTrue(torch.equal(original_output, calib_output))

        quant_model = quantizer.quantize(model)
        self.assertTrue(isinstance(quant_model.conv, torch.nn.Conv2d))
        self.assertTrue(isinstance(quant_model.sub.conv, torch.nn.Conv2d))
        self.assertEqual(quant_model.conv, model.conv)
        self.assertEqual(quant_model.sub.conv, model.sub.conv)
        self.assertNotEqual(quant_model.linear, model.linear)

        quant_output = quant_model(dummy_input)
        self.assertFalse(torch.equal(original_output, quant_output))
        torch.testing.assert_close(quant_output, original_output, rtol=0.1, atol=0.5)
//  [python]
def test_calib_amp_quantize(self, backend: Backend) -> None:
        model = SimpleModule()
        dummy_input = torch.randn((1, 2, 5, 5))
        quantizer = Quantizer(backend=backend)
        original_output = model(dummy_input)

        calib_model = quantizer.calib(model)
        calib_output = calib_model(dummy_input)
        self.assertTrue(torch.equal(original_output, calib_output))

        amp_model = quantizer.amp(model)
        amp_modules = dict(amp_model.named_modules())
        for name, mod in model.named_modules():
            if type(mod) in [torch.nn.Conv2d, torch.nn.Linear]:
                self.assertTrue(isinstance(amp_modules[name], AmpModule))
        amp_output = amp_model(dummy_input)
        self.assertTrue(torch.equal(original_output, amp_output))
        quantizer.fallback(amp_model, num=2)
        self.assertEqual(len(quantizer.module_filter.exclude_names), 2)

        quant_model = quantizer.quantize(model)
        torch.jit.trace(quant_model, dummy_input)
        modules = dict(model.named_modules())
        quant_modules = dict(quant_model.named_modules())
        for name in quantizer.module_filter.exclude_names:
            self.assertEqual(quant_modules[name], modules[name])
        quant_output = quant_model(dummy_input)
        self.assertFalse(torch.equal(original_output, quant_output))
        torch.testing.assert_allclose(quant_output, original_output, rtol=0.1, atol=0.5)
//  [python]
def test_calib_amp_quantize(self, backend: Backend) -> None:
        model = SimpleModule()
        dummy_input = torch.randn((1, 2, 5, 5))
        quantizer = Quantizer(backend=backend)
        original_output = model(dummy_input)

        calib_model = quantizer.calib(model)
        calib_output = calib_model(dummy_input)
        self.assertTrue(torch.equal(original_output, calib_output))

        amp_model = quantizer.amp(model)
        amp_modules = dict(amp_model.named_modules())
        for name, mod in model.named_modules():
            if type(mod) in [torch.nn.Conv2d, torch.nn.Linear]:
                self.assertTrue(isinstance(amp_modules[name], AmpModule))
        amp_output = amp_model(dummy_input)
        self.assertTrue(torch.equal(original_output, amp_output))
        quantizer.fallback(amp_model, num=2)
        self.assertEqual(len(quantizer.module_filter.exclude_names), 2)

        quant_model = quantizer.quantize(model)
        torch.jit.trace(quant_model, dummy_input)
        modules = dict(model.named_modules())
        quant_modules = dict(quant_model.named_modules())
        for name in quantizer.module_filter.exclude_names:
            self.assertEqual(quant_modules[name], modules[name])
        quant_output = quant_model(dummy_input)
        self.assertFalse(torch.equal(original_output, quant_output))
        torch.testing.assert_close(quant_output, original_output, rtol=0.1, atol=0.5)
//  [python]
def test_calib_and_quantize_with_module_filter(
        self, backend: Backend, module_filter: ModuleFilter
    ) -> None:
        model = UntraceableSimpleModule()
        quantizer = Quantizer(backend=backend, module_filter=module_filter)
        dummy_input = torch.randn((1, 2, 5, 5))
        original_output = model(dummy_input)

        calib_model = quantizer.calib(model)
        calib_output = calib_model(dummy_input)
        self.assertTrue(torch.equal(original_output, calib_output))

        qmodel = quantizer.quantize(model)
        self.assertTrue(isinstance(qmodel.traceable_sub.linear, torch.nn.Linear))
        self.assertTrue(isinstance(qmodel.traceable_sub.sub.conv, torch.nn.Conv2d))
        self.assertEqual(qmodel.traceable_sub.linear, model.traceable_sub.linear)
        self.assertEqual(qmodel.traceable_sub.sub.conv, model.traceable_sub.sub.conv)
        self.assertNotEqual(qmodel.traceable_sub.conv, model.traceable_sub.conv)

        quant_output = qmodel(dummy_input)
        torch.testing.assert_allclose(quant_output, original_output, rtol=0.1, atol=0.5)
//  [python]
def test_calib_and_quantize_with_module_filter(
        self, backend: Backend, module_filter: ModuleFilter
    ) -> None:
        model = UntraceableSimpleModule()
        quantizer = Quantizer(backend=backend, module_filter=module_filter)
        dummy_input = torch.randn((1, 2, 5, 5))
        original_output = model(dummy_input)

        calib_model = quantizer.calib(model)
        calib_output = calib_model(dummy_input)
        self.assertTrue(torch.equal(original_output, calib_output))

        qmodel = quantizer.quantize(model)
        self.assertTrue(isinstance(qmodel.traceable_sub.linear, torch.nn.Linear))
        self.assertTrue(isinstance(qmodel.traceable_sub.sub.conv, torch.nn.Conv2d))
        self.assertEqual(qmodel.traceable_sub.linear, model.traceable_sub.linear)
        self.assertEqual(qmodel.traceable_sub.sub.conv, model.traceable_sub.sub.conv)
        self.assertNotEqual(qmodel.traceable_sub.conv, model.traceable_sub.conv)

        quant_output = qmodel(dummy_input)
        torch.testing.assert_close(quant_output, original_output, rtol=0.1, atol=0.5)
//  [python]
model = model.cuda().eval()

# uncomment the following line to enable kernel tuning
# hidet.torch.dynamo_config.search_space(2)

# optimize the model with 'hidet' backend
model_opt = torch.compile(model, backend='hidet')

# run the optimized model
y1 = model_opt(x)
y2 = model(x)

# check the correctness
torch.testing.assert_allclose(actual=y1, expected=y2, rtol=1e-2, atol=1e-2)
//  [python]
model = model.cuda().eval()

# uncomment the following line to enable kernel tuning
# hidet.torch.dynamo_config.search_space(2)

# optimize the model with 'hidet' backend
model_opt = torch.compile(model, backend='hidet')

# run the optimized model
y1 = model_opt(x)
y2 = model(x)

# check the correctness
torch.testing.assert_close(actual=y1, expected=y2, rtol=1e-2, atol=1e-2)
//  [python]
def _test_against_cpu(self, inputs, reference, eps, dtype, is_causal=None):
        q, k, v, rpb, kernel_size, dilation = inputs
        kernel_size, dilation, is_causal = check_args(kernel_size, dilation, is_causal)
        attn_ref, out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            attn = na3d_qk(
                q_,
                k_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )
            attn = attn.softmax(dim=-1)
            out = na3d_av(
                attn,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_allclose(attn.float(), attn_ref, atol=eps, rtol=0)
            torch.testing.assert_allclose(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def _test_against_cpu(self, inputs, reference, eps, dtype, is_causal=None):
        q, k, v, rpb, kernel_size, dilation = inputs
        kernel_size, dilation, is_causal = check_args(kernel_size, dilation, is_causal)
        attn_ref, out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            attn = na3d_qk(
                q_,
                k_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )
            attn = attn.softmax(dim=-1)
            out = na3d_av(
                attn,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_close(attn.float(), attn_ref, atol=eps, rtol=0)
            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def test_fp8_dot_acc(in_type_str, low_precision_acc, device):
    check_type_supported(in_type_str, device)
    M, N, K = 128, 256, 256
    BLOCK_M, BLOCK_N, BLOCK_K = 128, 256, 128
    A = numpy_random((M, K), dtype_str=in_type_str)
    B = numpy_random((K, N), dtype_str=in_type_str)
    Bt = B.T
    C = torch.empty((M, N), dtype=torch.float32, device='cuda')
    num_warps = 8
    a = to_triton(A, device='cuda', dst_type=in_type_str)
    b = to_triton(B, device='cuda', dst_type=in_type_str)
    grid = (triton.cdiv(M, BLOCK_M), 1)
    matmul_kernel[grid](a, b, C, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), C.stride(0), C.stride(1),
                        BLOCK_M, BLOCK_N, BLOCK_K, low_precision_acc, num_warps=num_warps)
    torch_a = torch.from_numpy(A)
    th_a = f8_to_f16(torch_a.cuda(), in_type_str)
    torch_b = torch.from_numpy(B)
    th_b = f8_to_f16(torch_b.cuda(), in_type_str)
    ref_out = torch.matmul(th_a, th_b).to(torch.float32)
    if in_type_str == 'float8e4nv':
        torch.testing.assert_allclose(ref_out, C, rtol=0.01, atol=0.01)
    elif low_precision_acc > 32:
        torch.testing.assert_allclose(ref_out, C, rtol=1e-3, atol=1e-3)
    else:
        torch.testing.assert_allclose(ref_out, C)
//  [python]
def test_fp8_dot_acc(in_type_str, low_precision_acc, device):
    check_type_supported(in_type_str, device)
    M, N, K = 128, 256, 256
    BLOCK_M, BLOCK_N, BLOCK_K = 128, 256, 128
    A = numpy_random((M, K), dtype_str=in_type_str)
    B = numpy_random((K, N), dtype_str=in_type_str)
    Bt = B.T
    C = torch.empty((M, N), dtype=torch.float32, device='cuda')
    num_warps = 8
    a = to_triton(A, device='cuda', dst_type=in_type_str)
    b = to_triton(B, device='cuda', dst_type=in_type_str)
    grid = (triton.cdiv(M, BLOCK_M), 1)
    matmul_kernel[grid](a, b, C, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), C.stride(0), C.stride(1),
                        BLOCK_M, BLOCK_N, BLOCK_K, low_precision_acc, num_warps=num_warps)
    torch_a = torch.from_numpy(A)
    th_a = f8_to_f16(torch_a.cuda(), in_type_str)
    torch_b = torch.from_numpy(B)
    th_b = f8_to_f16(torch_b.cuda(), in_type_str)
    ref_out = torch.matmul(th_a, th_b).to(torch.float32)
    if in_type_str == 'float8e4nv':
        torch.testing.assert_close(ref_out, C, rtol=0.01, atol=0.01)
    elif low_precision_acc > 32:
        torch.testing.assert_close(ref_out, C, rtol=1e-3, atol=1e-3)
    else:
        torch.testing.assert_close(ref_out, C)
//  [python]
def test_max_value_writer_serialize(
        self, size, batch_size, reward_ranges, device, tmpdir
    ):
        rb = TensorDictReplayBuffer(
            storage=LazyTensorStorage(size, device=device),
            sampler=SamplerWithoutReplacement(),
            batch_size=batch_size,
            writer=TensorDictMaxValueWriter(rank_key="key"),
        )

        max_reward1, max_reward2, max_reward3 = reward_ranges

        td = TensorDict(
            {
                "key": torch.clamp_max(torch.rand(size), max=max_reward1),
                "obs": torch.rand(size),
            },
            batch_size=size,
            device=device,
        )
        rb.extend(td)
        rb._writer.dumps(tmpdir)
        # check we can dump twice
        rb._writer.dumps(tmpdir)
        other = TensorDictMaxValueWriter(rank_key="key")
        other.loads(tmpdir)
        assert len(rb._writer._current_top_values) == len(other._current_top_values)
        torch.testing.assert_allclose(
            torch.tensor(rb._writer._current_top_values),
            torch.tensor(other._current_top_values),
        )
//  [python]
def test_max_value_writer_serialize(
        self, size, batch_size, reward_ranges, device, tmpdir
    ):
        rb = TensorDictReplayBuffer(
            storage=LazyTensorStorage(size, device=device),
            sampler=SamplerWithoutReplacement(),
            batch_size=batch_size,
            writer=TensorDictMaxValueWriter(rank_key="key"),
        )

        max_reward1, max_reward2, max_reward3 = reward_ranges

        td = TensorDict(
            {
                "key": torch.clamp_max(torch.rand(size), max=max_reward1),
                "obs": torch.rand(size),
            },
            batch_size=size,
            device=device,
        )
        rb.extend(td)
        rb._writer.dumps(tmpdir)
        # check we can dump twice
        rb._writer.dumps(tmpdir)
        other = TensorDictMaxValueWriter(rank_key="key")
        other.loads(tmpdir)
        assert len(rb._writer._current_top_values) == len(other._current_top_values)
        torch.testing.assert_close(
            torch.tensor(rb._writer._current_top_values),
            torch.tensor(other._current_top_values),
        )
//  [python]
def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (3, 7, 5), (3, 2, 1)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        query = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 9, 14, 16, 16),
                torch.randn(2, 8, 9, 16, 18, 32),
                torch.randn(4, 1, 9, 32, 20, 16),
            ],
            **kwargs,
        )
        key = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 9, 14, 16, 16),
                torch.randn(2, 8, 9, 16, 18, 32),
                torch.randn(4, 1, 9, 32, 20, 16),
            ],
            **kwargs,
        )
        out_nested = na3d_qk(query, key, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for q, k in zip(query, key):
            out_ref.append(na3d_qk(q, k, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            t.assert_close(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            additional_keys = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na3d_qk(
                query,
                key,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_keys=additional_keys,
            )
            out_ref = []
            for q, k, k_add in zip(query, key, additional_keys):
                out_ref.append(
                    na3d_qk(
                        q,
                        k,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_keys=k_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                t.assert_close(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (3, 7, 5), (3, 2, 1)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        query = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 9, 14, 16, 16),
                torch.randn(2, 8, 9, 16, 18, 32),
                torch.randn(4, 1, 9, 32, 20, 16),
            ],
            **kwargs,
        )
        key = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 9, 14, 16, 16),
                torch.randn(2, 8, 9, 16, 18, 32),
                torch.randn(4, 1, 9, 32, 20, 16),
            ],
            **kwargs,
        )
        out_nested = na3d_qk(query, key, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for q, k in zip(query, key):
            out_ref.append(na3d_qk(q, k, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            additional_keys = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na3d_qk(
                query,
                key,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_keys=additional_keys,
            )
            out_ref = []
            for q, k, k_add in zip(query, key, additional_keys):
                out_ref.append(
                    na3d_qk(
                        q,
                        k,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_keys=k_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def test_generate(max_seq_length):
    T = 5
    input_idx = torch.randint(10, size=(T,))

    config = Config(block_size=128, vocab_size=16, n_layer=1, n_head=4, n_embd=8)
    model = GPT(config)
    model.max_seq_length = max_seq_length
    model.set_kv_cache(batch_size=1)
    max_new_tokens = 20

    multinomial_results = []

    def multinomial(*args, **kwargs):
        out = torch.multinomial(*args, **kwargs, num_samples=1)
        multinomial_results.append(out)
        return out

    with mock.patch("litgpt.generate.base.multinomial_num_samples_1", multinomial):
        out = generate.generate(model, input_idx, T + max_new_tokens, top_k=4)

    assert out.size(0) == T + max_new_tokens
    multinomial_results = torch.hstack(multinomial_results)
    expected = torch.cat((input_idx, multinomial_results))
    assert out.shape == expected.shape
    torch.testing.assert_allclose(out, expected)
//  [python]
def test_generate(max_seq_length):
    T = 5
    input_idx = torch.randint(10, size=(T,))

    config = Config(block_size=128, vocab_size=16, n_layer=1, n_head=4, n_embd=8)
    model = GPT(config)
    model.max_seq_length = max_seq_length
    model.set_kv_cache(batch_size=1)
    max_new_tokens = 20

    multinomial_results = []

    def multinomial(*args, **kwargs):
        out = torch.multinomial(*args, **kwargs, num_samples=1)
        multinomial_results.append(out)
        return out

    with mock.patch("litgpt.generate.base.multinomial_num_samples_1", multinomial):
        out = generate.generate(model, input_idx, T + max_new_tokens, top_k=4)

    assert out.size(0) == T + max_new_tokens
    multinomial_results = torch.hstack(multinomial_results)
    expected = torch.cat((input_idx, multinomial_results))
    assert out.shape == expected.shape
    torch.testing.assert_close(out, expected)
//  [python]
def test_stft_l2():
    N = 32
    n = torch.arange(N)

    f = N / 4
    target = torch.cos(2 * math.pi * f * n / N)
    target = target[None, None, :]
    pred = torch.zeros_like(target)

    loss = auraloss.freq.STFTLoss(
        fft_size=N,
        hop_size=N + 1,  # eliminate padding artefacts by enforcing only one hop
        win_length=N,
        window="ones",  # eliminate windowing artefacts
        w_sc=0.0,
        w_log_mag=0.0,
        w_lin_mag=1.0,
        w_phs=0.0,
        mag_distance="L2",
    )
    res = loss(pred, target)

    # MSE of energy in concentrated in single DFT bin
    expected_loss = ((N // 2) ** 2) / (N // 2 + 1)

    torch.testing.assert_allclose(res, torch.tensor(expected_loss), rtol=1e-3, atol=1e-3)
//  [python]
def test_stft_l2():
    N = 32
    n = torch.arange(N)

    f = N / 4
    target = torch.cos(2 * math.pi * f * n / N)
    target = target[None, None, :]
    pred = torch.zeros_like(target)

    loss = auraloss.freq.STFTLoss(
        fft_size=N,
        hop_size=N + 1,  # eliminate padding artefacts by enforcing only one hop
        win_length=N,
        window="ones",  # eliminate windowing artefacts
        w_sc=0.0,
        w_log_mag=0.0,
        w_lin_mag=1.0,
        w_phs=0.0,
        mag_distance="L2",
    )
    res = loss(pred, target)

    # MSE of energy in concentrated in single DFT bin
    expected_loss = ((N // 2) ** 2) / (N // 2 + 1)

    torch.testing.assert_close(res, torch.tensor(expected_loss), rtol=1e-3, atol=1e-3)
//  [python]
def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (3, 7, 5), (3, 2, 1)
        num_na_weights = get_num_na_weights(kernel_size)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        attn = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 9, 14, 16, num_na_weights),
                torch.randn(2, 8, 9, 16, 18, num_na_weights),
                torch.randn(4, 1, 9, 32, 20, num_na_weights),
            ],
            **kwargs,
        )
        value = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 9, 14, 16, 16),
                torch.randn(2, 8, 9, 16, 18, 32),
                torch.randn(4, 1, 9, 32, 20, 16),
            ],
            **kwargs,
        )
        out_nested = na3d_av(attn, value, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for a, v in zip(attn, value):
            out_ref.append(na3d_av(a, v, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_allclose(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            attn = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, 9, 14, 16, num_na_weights + n_extra_tokens),
                    torch.randn(2, 8, 9, 16, 18, num_na_weights + n_extra_tokens),
                    torch.randn(4, 1, 9, 32, 20, num_na_weights + n_extra_tokens),
                ],
                **kwargs,
            )
            additional_values = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na3d_av(
                attn,
                value,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_values=additional_values,
            )
            out_ref = []
            for a, v, v_add in zip(attn, value, additional_values):
                out_ref.append(
                    na3d_av(
                        a,
                        v,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_values=v_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_allclose(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (3, 7, 5), (3, 2, 1)
        num_na_weights = get_num_na_weights(kernel_size)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        attn = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 9, 14, 16, num_na_weights),
                torch.randn(2, 8, 9, 16, 18, num_na_weights),
                torch.randn(4, 1, 9, 32, 20, num_na_weights),
            ],
            **kwargs,
        )
        value = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 9, 14, 16, 16),
                torch.randn(2, 8, 9, 16, 18, 32),
                torch.randn(4, 1, 9, 32, 20, 16),
            ],
            **kwargs,
        )
        out_nested = na3d_av(attn, value, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for a, v in zip(attn, value):
            out_ref.append(na3d_av(a, v, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            attn = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, 9, 14, 16, num_na_weights + n_extra_tokens),
                    torch.randn(2, 8, 9, 16, 18, num_na_weights + n_extra_tokens),
                    torch.randn(4, 1, 9, 32, 20, num_na_weights + n_extra_tokens),
                ],
                **kwargs,
            )
            additional_values = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na3d_av(
                attn,
                value,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_values=additional_values,
            )
            out_ref = []
            for a, v, v_add in zip(attn, value, additional_values):
                out_ref.append(
                    na3d_av(
                        a,
                        v,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_values=v_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def test_multires_l2():
    N = 32
    n = torch.arange(N)

    f = N / 4
    target = torch.cos(2 * math.pi * f * n / N)
    target = target[None, None, :]
    pred = torch.zeros_like(target)

    loss = auraloss.freq.MultiResolutionSTFTLoss(
        fft_sizes=[N],
        hop_sizes=[N + 1],  # eliminate padding artefacts by enforcing only one hop
        win_lengths=[N],
        window="ones",  # eliminate windowing artefacts
        w_sc=0.0,
        w_log_mag=0.0,
        w_lin_mag=1.0,
        w_phs=0.0,
        mag_distance="L2",
    )
    res = loss(pred, target)

    expected_loss = ((N // 2) ** 2) / (N // 2 + 1)

    torch.testing.assert_allclose(res, torch.tensor(expected_loss), rtol=1e-3, atol=1e-3)
//  [python]
def test_multires_l2():
    N = 32
    n = torch.arange(N)

    f = N / 4
    target = torch.cos(2 * math.pi * f * n / N)
    target = target[None, None, :]
    pred = torch.zeros_like(target)

    loss = auraloss.freq.MultiResolutionSTFTLoss(
        fft_sizes=[N],
        hop_sizes=[N + 1],  # eliminate padding artefacts by enforcing only one hop
        win_lengths=[N],
        window="ones",  # eliminate windowing artefacts
        w_sc=0.0,
        w_log_mag=0.0,
        w_lin_mag=1.0,
        w_phs=0.0,
        mag_distance="L2",
    )
    res = loss(pred, target)

    expected_loss = ((N // 2) ** 2) / (N // 2 + 1)

    torch.testing.assert_close(res, torch.tensor(expected_loss), rtol=1e-3, atol=1e-3)
//  [python]
def test_label_smoothing(self):
        pad_index = 0
        smoothing = 0.4
        criterion = XentLoss(pad_index=pad_index, smoothing=smoothing)

        # batch x seq_len x vocab_size: 3 x 2 x 5
        predict = torch.FloatTensor([
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
        ])

        # batch x seq_len: 3 x 2
        targets = torch.LongTensor([[2, 1], [2, 0], [1, 0]])

        # test the smoothing function
        # pylint: disable=protected-access
        smoothed_targets = criterion._smooth_targets(
            targets=targets.view(-1), vocab_size=predict.size(-1)
        )
        # pylint: enable=protected-access
        torch.testing.assert_allclose(
            smoothed_targets,
            torch.Tensor([
                [0.0000, 0.1333, 0.6000, 0.1333, 0.1333],
                [0.0000, 0.6000, 0.1333, 0.1333, 0.1333],
                [0.0000, 0.1333, 0.6000, 0.1333, 0.1333],
                [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                [0.0000, 0.6000, 0.1333, 0.1333, 0.1333],
                [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
            ]),
            rtol=1e-4,
            atol=1e-4,
        )
        self.assertEqual(torch.max(smoothed_targets), 1 - smoothing)

        # test the loss computation
        v = criterion(predict.log(), **{"trg": targets})
        self.assertAlmostEqual(v.item(), 2.1326, places=4)
//  [python]
def test_label_smoothing(self):
        pad_index = 0
        smoothing = 0.4
        criterion = XentLoss(pad_index=pad_index, smoothing=smoothing)

        # batch x seq_len x vocab_size: 3 x 2 x 5
        predict = torch.FloatTensor([
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
        ])

        # batch x seq_len: 3 x 2
        targets = torch.LongTensor([[2, 1], [2, 0], [1, 0]])

        # test the smoothing function
        # pylint: disable=protected-access
        smoothed_targets = criterion._smooth_targets(
            targets=targets.view(-1), vocab_size=predict.size(-1)
        )
        # pylint: enable=protected-access
        torch.testing.assert_close(
            smoothed_targets,
            torch.Tensor([
                [0.0000, 0.1333, 0.6000, 0.1333, 0.1333],
                [0.0000, 0.6000, 0.1333, 0.1333, 0.1333],
                [0.0000, 0.1333, 0.6000, 0.1333, 0.1333],
                [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                [0.0000, 0.6000, 0.1333, 0.1333, 0.1333],
                [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
            ]),
            rtol=1e-4,
            atol=1e-4,
        )
        self.assertEqual(torch.max(smoothed_targets), 1 - smoothing)

        # test the loss computation
        v = criterion(predict.log(), **{"trg": targets})
        self.assertAlmostEqual(v.item(), 2.1326, places=4)
//  [python]
def test_no_label_smoothing(self):
        pad_index = 0
        smoothing = 0.0
        criterion = XentLoss(pad_index=pad_index, smoothing=smoothing)

        # batch x seq_len x vocab_size: 3 x 2 x 5
        predict = torch.FloatTensor([
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
        ])

        # batch x seq_len: 3 x 2
        targets = torch.LongTensor([[2, 1], [2, 0], [1, 0]])

        # test the smoothing function: should still be one-hot
        # pylint: disable=protected-access
        smoothed_targets = criterion._smooth_targets(
            targets=targets.view(-1), vocab_size=predict.size(-1)
        )
        # pylint: enable=protected-access

        self.assertEqual(torch.max(smoothed_targets), 1)
        self.assertEqual(torch.min(smoothed_targets), 0)

        self.assertTrue(torch.allclose(
            smoothed_targets,
            torch.Tensor([
                [0.0, 0.0, 1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
            ]),
            rtol=1e-4,
            atol=1e-4,
        ))

        v = criterion(predict.log(), **{"trg": targets})
        self.assertAlmostEqual(v.item(), 5.6268, places=4)
//  [python]
def test_no_label_smoothing(self):
        pad_index = 0
        smoothing = 0.0
        criterion = XentLoss(pad_index=pad_index, smoothing=smoothing)

        # batch x seq_len x vocab_size: 3 x 2 x 5
        predict = torch.FloatTensor([
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
        ])

        # batch x seq_len: 3 x 2
        targets = torch.LongTensor([[2, 1], [2, 0], [1, 0]])

        # test the smoothing function: should still be one-hot
        # pylint: disable=protected-access
        smoothed_targets = criterion._smooth_targets(
            targets=targets.view(-1), vocab_size=predict.size(-1)
        )
        # pylint: enable=protected-access

        self.assertEqual(torch.max(smoothed_targets), 1)
        self.assertEqual(torch.min(smoothed_targets), 0)

        torch.testing.assert_close(
            smoothed_targets,
            torch.Tensor([
                [0.0, 0.0, 1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
            ]),
            rtol=1e-4,
            atol=1e-4,
        )

        v = criterion(predict.log(), **{"trg": targets})
        self.assertAlmostEqual(v.item(), 5.6268, places=4)
//  [python]
def _test_softmax(
        self,
        batch_sizes=(1, 1024),
        input_shapes=(6,),
        dim=-1,
        dtype="float16",
        testname="softmax",
    ):
        module = self._build_model(batch_sizes, input_shapes, dim, dtype, testname)
        torch_dtype = string_to_torch_dtype(dtype)

        for batch_size in batch_sizes:
            x_pt = torch.randn(batch_size, *input_shapes, dtype=torch_dtype).cuda()
            y_pt = torch.nn.functional.softmax(x_pt, dim=dim)

            y = torch.empty([batch_size, *input_shapes], dtype=torch_dtype).cuda()
            module.run_with_tensors([x_pt], [y])
            torch.testing.assert_allclose(y_pt, y, atol=1e-2, rtol=1e-2)
//  [python]
def _test_softmax(
        self,
        batch_sizes=(1, 1024),
        input_shapes=(6,),
        dim=-1,
        dtype="float16",
        testname="softmax",
    ):
        module = self._build_model(batch_sizes, input_shapes, dim, dtype, testname)
        torch_dtype = string_to_torch_dtype(dtype)

        for batch_size in batch_sizes:
            x_pt = torch.randn(batch_size, *input_shapes, dtype=torch_dtype).cuda()
            y_pt = torch.nn.functional.softmax(x_pt, dim=dim)

            y = torch.empty([batch_size, *input_shapes], dtype=torch_dtype).cuda()
            module.run_with_tensors([x_pt], [y])
            torch.testing.assert_close(y_pt, y, atol=1e-2, rtol=1e-2)
//  [python]
def test_expand(rank, scale):
    pod = PODBlock(rank, scale)
    pod.fit(toy_snapshots)
    c = pod(toy_snapshots)
    t.assert_close(pod.expand(c), toy_snapshots)
    t.assert_close(pod.expand(c[0]), toy_snapshots[0].unsqueeze(0))
//  [python]
def test_expand(rank, scale):
    pod = PODBlock(rank, scale)
    pod.fit(toy_snapshots)
    c = pod(toy_snapshots)
    torch.testing.assert_close(pod.expand(c), toy_snapshots)
    torch.testing.assert_close(pod.expand(c[0]), toy_snapshots[0].unsqueeze(0))
//  [python]
def test_reduce_expand(rank, scale):
    pod = PODBlock(rank, scale)
    pod.fit(toy_snapshots)
    t.assert_close(
        pod.expand(pod.reduce(toy_snapshots)),
        toy_snapshots)
    t.assert_close(
        pod.expand(pod.reduce(toy_snapshots[0])),
        toy_snapshots[0].unsqueeze(0))
//  [python]
def test_reduce_expand(rank, scale):
    pod = PODBlock(rank, scale)
    pod.fit(toy_snapshots)
    torch.testing.assert_close(
        pod.expand(pod.reduce(toy_snapshots)),
        toy_snapshots)
    torch.testing.assert_close(
        pod.expand(pod.reduce(toy_snapshots[0])),
        toy_snapshots[0].unsqueeze(0))
//  [python]
def test_conditioning(self):
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()

        labels = self.prepare_inputs()
        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long) for _ in range(3)]

        top_prior = model.priors[0]
        start = 0
        music_token_conds = top_prior.get_music_tokens_conds(zs, start=start, end=start + top_prior.n_ctx)
        metadata = top_prior.get_metadata(labels[0].clone(), start, 1058304, 0)

        self.assertIsNone(music_token_conds)
        self.assertListEqual(metadata.numpy()[0][:10].tolist(), self.EXPECTED_Y_COND)

        audio_conditioning, metadata_conditioning, lyric_tokens = top_prior.get_cond(music_token_conds, metadata)
        self.assertTrue(torch.allclose(
            audio_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_AUDIO_COND), atol=1e-4, rtol=1e-4
        ))
        self.assertTrue(torch.allclose(
            metadata_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_META_COND), atol=1e-4, rtol=1e-4
        ))
        self.assertTrue(torch.allclose(
            lyric_tokens[0, :30].detach(), torch.tensor(self.EXPECTED_LYRIC_COND), atol=1e-4, rtol=1e-4
        ))
//  [python]
def test_forward():
    pod = PODBlock(1)
    pod.fit(toy_snapshots)
    c = pod(toy_snapshots)
    assert c.shape[0] == toy_snapshots.shape[0]
    assert c.shape[1] == pod.rank
    t.assert_close(c.mean(dim=0), torch.zeros(pod.rank))
    t.assert_close(c.std(dim=0), torch.ones(pod.rank))

    c = pod(toy_snapshots[0])
    assert c.shape[1] == pod.rank
    assert c.shape[0] == 1

    pod = PODBlock(2, False)
    pod.fit(toy_snapshots)
    c = pod(toy_snapshots)
    t.assert_close(c, (pod.basis @ toy_snapshots.T).T)
    with pytest.raises(AssertionError):
        t.assert_close(c.mean(dim=0), torch.zeros(pod.rank))
        t.assert_close(c.std(dim=0), torch.ones(pod.rank))
//  [python]
def test_forward():
    pod = PODBlock(1)
    pod.fit(toy_snapshots)
    c = pod(toy_snapshots)
    assert c.shape[0] == toy_snapshots.shape[0]
    assert c.shape[1] == pod.rank
    torch.testing.assert_close(c.mean(dim=0), torch.zeros(pod.rank))
    torch.testing.assert_close(c.std(dim=0), torch.ones(pod.rank))

    c = pod(toy_snapshots[0])
    assert c.shape[1] == pod.rank
    assert c.shape[0] == 1

    pod = PODBlock(2, False)
    pod.fit(toy_snapshots)
    c = pod(toy_snapshots)
    torch.testing.assert_close(c, (pod.basis @ toy_snapshots.T).T)
    with pytest.raises(AssertionError):
        torch.testing.assert_close(c.mean(dim=0), torch.zeros(pod.rank))
        torch.testing.assert_close(c.std(dim=0), torch.ones(pod.rank))
//  [python]
def test_SpectralConv3D(factorization, implementation):
    """Compare generic SpectralConv with hand written SpectralConv2D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    Note that this implies the order in which the conv is done in the manual implementation matches the automatic one, 
    take with a grain of salt
    """
    conv = SpectralConv(
        3, 6, (4, 5, 2), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )

    conv_dense = SpectralConv3d(
        3, 6, (4, 5, 2), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 3, 12, 12, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    t.assert_close(res_dense, res)
//  [python]
def test_SpectralConv3D(factorization, implementation):
    """Compare generic SpectralConv with hand written SpectralConv2D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    Note that this implies the order in which the conv is done in the manual implementation matches the automatic one, 
    take with a grain of salt
    """
    conv = SpectralConv(
        3, 6, (4, 5, 2), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )

    conv_dense = SpectralConv3d(
        3, 6, (4, 5, 2), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 3, 12, 12, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    torch.testing.assert_close(res_dense, res)
//  [python]
def test_SpectralConv1D(factorization, implementation):
    """Test for SpectralConv1D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    """
    conv = SpectralConv(
        10, 11, (5,), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )

    conv_dense = SpectralConv1d(
        10, 11, (5,), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 10, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    torch.allclose(res_dense, res)