jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_trace_buffer_slice(self):
        bs, d_hid = 10, 23

        class ExampleCode(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.mm_param = torch.nn.Parameter(torch.randn(d_hid, d_hid))
                self.mm_param2 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
                self.lin = torch.nn.Linear(d_hid, d_hid)
                self.register_buffer('buffer', torch.randn(bs + 100, d_hid))

            def forward(self, x):
                x = torch.mm(x, self.mm_param)
                skip_connection = x
                x = torch.relu(x)
                x = torch.mm(x, self.mm_param) + self.buffer[:x.shape[0]]
                x = self.lin(x)
                x = torch.relu(x)
                x = x + skip_connection
                x = torch.mm(x, self.mm_param2)
                x = self.lin(x)
                return x

        ec = ExampleCode()

        traced = torch.fx.symbolic_trace(ec)

        x = torch.randn(bs, d_hid)
        torch.testing.assert_close(ec(x), traced(x))
//  [python]
def test_frequency_with_valid_input(self) -> None:
        input = torch.tensor(
            [0.4826, 0.9517, 0.8967, 0.8995, 0.1584, 0.9445, 0.9700],
        )

        t.assert_allclose(
            frequency_at_k(input, k=0.5),
            torch.tensor([1.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000]),
        )
        t.assert_allclose(
            frequency_at_k(input, k=0.9),
            torch.tensor([1.0000, 0.0000, 1.0000, 1.0000, 1.0000, 0.0000, 0.0000]),
        )
        t.assert_allclose(
            frequency_at_k(input, k=0.95),
            torch.tensor([1.0000, 0.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.0000]),
        )
        t.assert_allclose(
            frequency_at_k(input, k=1.0),
            torch.tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]),
        )
//  [python]
def test_frequency_with_valid_input(self) -> None:
        input = torch.tensor(
            [0.4826, 0.9517, 0.8967, 0.8995, 0.1584, 0.9445, 0.9700],
        )

        torch.testing.assert_close(
            frequency_at_k(input, k=0.5),
            torch.tensor([1.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000]),
        )
        torch.testing.assert_close(
            frequency_at_k(input, k=0.9),
            torch.tensor([1.0000, 0.0000, 1.0000, 1.0000, 1.0000, 0.0000, 0.0000]),
        )
        torch.testing.assert_close(
            frequency_at_k(input, k=0.95),
            torch.tensor([1.0000, 0.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.0000]),
        )
        torch.testing.assert_close(
            frequency_at_k(input, k=1.0),
            torch.tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]),
        )
//  [python]
def test_call_function_ex(jit):
    jitting = False

    def foo(a, b):
        assert is_jitting_with_raise() == jitting
        return a + b

    def argsplat(*args):
        assert is_jitting_with_raise() == jitting
        return foo(*args)

    def kwargsplat(**kwargs):
        assert is_jitting_with_raise() == jitting
        return foo(**kwargs)

    assert any(i.opname == "CALL_FUNCTION_EX" and not i.arg & 1 for i in dis.get_instructions(argsplat))
    assert any(i.opname == "CALL_FUNCTION_EX" and i.arg & 1 for i in dis.get_instructions(kwargsplat))

    kwargs = {"a": 1, "b": 2}

    jitting = False
    res1 = argsplat(*kwargs.values())
    res2 = kwargsplat(**kwargs)

    jitting = True
    jres1 = jit(argsplat)(*kwargs.values())
    jres2 = jit(kwargsplat)(**kwargs)

    t.assert_close(res1, jres1)
    t.assert_close(res2, jres2)
//  [python]
def test_call_function_ex(jit):
    jitting = False

    def foo(a, b):
        assert is_jitting_with_raise() == jitting
        return a + b

    def argsplat(*args):
        assert is_jitting_with_raise() == jitting
        return foo(*args)

    def kwargsplat(**kwargs):
        assert is_jitting_with_raise() == jitting
        return foo(**kwargs)

    assert any(i.opname == "CALL_FUNCTION_EX" and not i.arg & 1 for i in dis.get_instructions(argsplat))
    assert any(i.opname == "CALL_FUNCTION_EX" and i.arg & 1 for i in dis.get_instructions(kwargsplat))

    kwargs = {"a": 1, "b": 2}

    jitting = False
    res1 = argsplat(*kwargs.values())
    res2 = kwargsplat(**kwargs)

    jitting = True
    jres1 = jit(argsplat)(*kwargs.values())
    jres2 = jit(kwargsplat)(**kwargs)

    assert_close(res1, jres1)
    assert_close(res2, jres2)
//  [python]
def _compare_onnx_pytorch_outputs_in_np(
    onnx_outs: _OutputsType,
    pt_outs: _OutputsType,
    options: VerificationOptions,
):
    assert len(onnx_outs) == len(
        pt_outs
    ), f"Number of outputs differ ONNX runtime: ({len(onnx_outs)}) PyTorch: ({len(pt_outs)})"
    acceptable_error_percentage = options.acceptable_error_percentage
    if acceptable_error_percentage and (
        acceptable_error_percentage > 1.0 or acceptable_error_percentage < 0.0
    ):
        raise ValueError(
            "If set, acceptable_error_percentage should be between 0.0 and 1.0"
        )

    for ort_out, pt_out in zip(onnx_outs, pt_outs):
        try:
            # TODO: Remove `check_shape` option once every shape inconsistent issue is addressed.
            if not options.check_shape:
                # Allow different but broadcastable output shapes.
                ort_out, pt_out = np.broadcast_arrays(ort_out, pt_out)
            t.assert_allclose(
                ort_out,
                pt_out,
                rtol=options.rtol,
                atol=options.atol,
                check_dtype=options.check_dtype,
                equal_nan=True,
            )
        except AssertionError as e:
            if acceptable_error_percentage:
                error_percentage = 1 - np.sum(
                    np.isclose(ort_out, pt_out, rtol=options.rtol, atol=options.atol)
                ) / np.prod(ort_out.shape)
                if error_percentage <= acceptable_error_percentage:
                    warnings.warn(
                        f"Suppressed AssertionError:\n{e}.\n"
                        f"Error percentage {error_percentage} "
                        f"within acceptable range {acceptable_error_percentage}."
                    )
                    continue
            if ort_out.dtype == np.uint8 or ort_out.dtype == np.int8:
                warnings.warn("ONNX output is quantized")
            if pt_out.dtype == np.uint8 or pt_out.dtype == np.int8:
                warnings.warn("PyTorch output is quantized")
            raise
//  [python]
def test_build_map_dict_merge(jit):
    addall = lambda *args, **kwargs: sum(args) + sum(kwargs.values())
    foo = lambda *args, **kwargs: addall(*args, **kwargs)

    assert any(i.opname == "BUILD_MAP" for i in dis.get_instructions(foo))
    assert any(i.opname == "DICT_MERGE" for i in dis.get_instructions(foo))

    jfoo = jit(foo)

    args = (4, 3)
    kwargs = {"a": 1, "b": 2}

    thunder_result = jfoo(*args, **kwargs)
    python_result = foo(*args, **kwargs)

    with pytest.raises(KeyError, match="got multiple values for keyword argument") as excinfo:
        d = {"a": 3, "b": 4}
        mergefail = lambda **kwargs: addall(**kwargs, **d)
        jfail = jit(mergefail)
        jfail(**kwargs)

    t.assert_close(thunder_result, python_result)
//  [python]
def test_build_map_dict_merge(jit):
    addall = lambda *args, **kwargs: sum(args) + sum(kwargs.values())
    foo = lambda *args, **kwargs: addall(*args, **kwargs)

    assert any(i.opname == "BUILD_MAP" for i in dis.get_instructions(foo))
    assert any(i.opname == "DICT_MERGE" for i in dis.get_instructions(foo))

    jfoo = jit(foo)

    args = (4, 3)
    kwargs = {"a": 1, "b": 2}

    thunder_result = jfoo(*args, **kwargs)
    python_result = foo(*args, **kwargs)

    with pytest.raises(KeyError, match="got multiple values for keyword argument") as excinfo:
        d = {"a": 3, "b": 4}
        mergefail = lambda **kwargs: addall(**kwargs, **d)
        jfail = jit(mergefail)
        jfail(**kwargs)

    assert_close(thunder_result, python_result)
//  [python]
def test_dict_update(jit):
    jitting = False

    def addall(*args, **kwargs):
        assert is_jitting_with_raise() == jitting
        return sum(args) + sum(kwargs.values())

    def foo(*args, **kwargs):
        assert is_jitting_with_raise() == jitting
        return addall(*args, **{**kwargs, "x": 1})

    assert any(i.opname == "DICT_UPDATE" for i in dis.get_instructions(foo))

    args = (4, 3)
    kwargs = {"a": 1, "b": 2}

    jitting = False
    python_result = foo(*args, **kwargs)
    jitting = True
    thunder_result = jit(foo)(*args, **kwargs)

    t.assert_close(thunder_result, python_result)
//  [python]
def test_dict_update(jit):
    jitting = False

    def addall(*args, **kwargs):
        assert is_jitting_with_raise() == jitting
        return sum(args) + sum(kwargs.values())

    def foo(*args, **kwargs):
        assert is_jitting_with_raise() == jitting
        return addall(*args, **{**kwargs, "x": 1})

    assert any(i.opname == "DICT_UPDATE" for i in dis.get_instructions(foo))

    args = (4, 3)
    kwargs = {"a": 1, "b": 2}

    jitting = False
    python_result = foo(*args, **kwargs)
    jitting = True
    thunder_result = jit(foo)(*args, **kwargs)

    assert_close(thunder_result, python_result)
//  [python]
def test_inner_closure(jit):
    # NOTE The addition of closing over value also tests
    #   the STORE_DEREF opcode
    def foo(a, b):
        value = 5

        def bar(a):
            return a + b + value

        return bar(a + 1)

    jfoo = jit(foo)

    args = (4, 3)

    thunder_result = jfoo(*args)
    python_result = foo(*args)

    t.assert_close(thunder_result, python_result)
//  [python]
def test_inner_closure(jit):
    # NOTE The addition of closing over value also tests
    #   the STORE_DEREF opcode
    def foo(a, b):
        value = 5

        def bar(a):
            return a + b + value

        return bar(a + 1)

    jfoo = jit(foo)

    args = (4, 3)

    thunder_result = jfoo(*args)
    python_result = foo(*args)

    assert_close(thunder_result, python_result)
//  [python]
def test_benchmark(self):
        split_ks = sorted(set(range(1, 6)).union([2**i for i in range(5)]))
        for split_k, (n, k) in itertools.product(split_ks, NK_SHAPES):
            NUM_ITERS = 100000
            NUM_WARMUP_ITERS = 1000
            ait_module = build_ait_module_gemm_rcr(
                ms=BATCH_SIZES,
                n=n,
                k=k,
                split_k=split_k,
                test_name=f"gemm_rcr_{split_k=}_{uuid.uuid4().hex}",
            )
            for m in BATCH_SIZES:
                mnk = {"m": m, "n": n, "k": k}
                _LOGGER.warning(f"mnk={mnk}, split_k={split_k}")
                inputs_pool, outputs_pool = make_input_output_pools(
                    pool_size=INPUT_POOL_SIZE,
                    eval_pt_func=lambda: eval_pt_gemm_rcr(**mnk),
                    input_filter_func=lambda name, _: not name.startswith("output"),
                    output_filter_func=lambda name, _: name.startswith("output"),
                )
                gemm_rcr_function = GemmRCRFunction(inputs_pool)
                gemm_rcr_trt_function = GemmRCRTRTFunction(
                    inputs_pool, max_batch_size=m
                )

                pt_outputs = eval_pt_gemm_rcr(**mnk)
                ait_outputs = {"output": torch.empty_like(pt_outputs["output"])}
                ait_module.run_with_tensors(
                    {k: v for k, v in pt_outputs.items() if k != "output"},
                    ait_outputs,
                )
                torch.testing.assert_allclose(
                    ait_outputs["output"], pt_outputs["output"], rtol=1, atol=1
                )
                mean_runtime_ait = run_benchmark(
                    ait_module=ait_module,
                    inputs_pool=inputs_pool,
                    outputs_pool=outputs_pool,
                    num_iters=NUM_ITERS,
                    num_warmup_iters=NUM_WARMUP_ITERS,
                )

                mean_runtime_pt = benchmark_torch_function(
                    iters=NUM_ITERS, function=gemm_rcr_function
                )

                mean_runtime_trt = benchmark_torch_function(
                    iters=NUM_ITERS, function=gemm_rcr_trt_function
                )

                benchmark_results = {
                    "function": "gemm_rcr_bias",
                    "mean_runtime_ait_ms": round(mean_runtime_ait, 5),
                    "mean_runtime_pt_ms": round(mean_runtime_pt, 5),
                    "mean_runtime_trt_ms": round(mean_runtime_trt, 5),
                    "split_k": split_k,
                    **mnk,
                }
                _LOGGER.warning(
                    f"Benchmark results {json.dumps(benchmark_results, separators=(',', ':'))}",
                )
//  [python]
def test_benchmark(self):
        split_ks = sorted(set(range(1, 6)).union([2**i for i in range(5)]))
        for split_k, (n, k) in itertools.product(split_ks, NK_SHAPES):
            NUM_ITERS = 100000
            NUM_WARMUP_ITERS = 1000
            ait_module = build_ait_module_gemm_rcr(
                ms=BATCH_SIZES,
                n=n,
                k=k,
                split_k=split_k,
                test_name=f"gemm_rcr_{split_k=}_{uuid.uuid4().hex}",
            )
            for m in BATCH_SIZES:
                mnk = {"m": m, "n": n, "k": k}
                _LOGGER.warning(f"mnk={mnk}, split_k={split_k}")
                inputs_pool, outputs_pool = make_input_output_pools(
                    pool_size=INPUT_POOL_SIZE,
                    eval_pt_func=lambda: eval_pt_gemm_rcr(**mnk),
                    input_filter_func=lambda name, _: not name.startswith("output"),
                    output_filter_func=lambda name, _: name.startswith("output"),
                )
                gemm_rcr_function = GemmRCRFunction(inputs_pool)
                gemm_rcr_trt_function = GemmRCRTRTFunction(
                    inputs_pool, max_batch_size=m
                )

                pt_outputs = eval_pt_gemm_rcr(**mnk)
                ait_outputs = {"output": torch.empty_like(pt_outputs["output"])}
                ait_module.run_with_tensors(
                    {k: v for k, v in pt_outputs.items() if k != "output"},
                    ait_outputs,
                )
                torch.testing.assert_close(
                    ait_outputs["output"], pt_outputs["output"], rtol=1, atol=1
                )
                mean_runtime_ait = run_benchmark(
                    ait_module=ait_module,
                    inputs_pool=inputs_pool,
                    outputs_pool=outputs_pool,
                    num_iters=NUM_ITERS,
                    num_warmup_iters=NUM_WARMUP_ITERS,
                )

                mean_runtime_pt = benchmark_torch_function(
                    iters=NUM_ITERS, function=gemm_rcr_function
                )

                mean_runtime_trt = benchmark_torch_function(
                    iters=NUM_ITERS, function=gemm_rcr_trt_function
                )

                benchmark_results = {
                    "function": "gemm_rcr_bias",
                    "mean_runtime_ait_ms": round(mean_runtime_ait, 5),
                    "mean_runtime_pt_ms": round(mean_runtime_pt, 5),
                    "mean_runtime_trt_ms": round(mean_runtime_trt, 5),
                    "split_k": split_k,
                    **mnk,
                }
                _LOGGER.warning(
                    f"Benchmark results {json.dumps(benchmark_results, separators=(',', ':'))}",
                )
//  [python]
def test_inner_function_definition(jit):
    def foo(a, b):
        def bar(a, b):
            return a + b

        return bar(a + 1, b)

    jfoo = jit(foo)

    args = (4, 3)

    thunder_result = jfoo(*args)
    python_result = foo(*args)

    t.assert_close(thunder_result, python_result)

    def foo(a, b):
        def bar(a, b=b):
            return a + b

        return bar(a + 1)

    t.assert_close(foo(*args), jit(foo)(*args))

    def foo(a, b):
        def bar(a: int, *, b: int = b):
            return a + b

        return bar(a + 1)

    t.assert_close(foo(*args), jit(foo)(*args))
//  [python]
def test_inner_function_definition(jit):
    def foo(a, b):
        def bar(a, b):
            return a + b

        return bar(a + 1, b)

    jfoo = jit(foo)

    args = (4, 3)

    thunder_result = jfoo(*args)
    python_result = foo(*args)

    assert_close(thunder_result, python_result)

    def foo(a, b):
        def bar(a, b=b):
            return a + b

        return bar(a + 1)

    assert_close(foo(*args), jit(foo)(*args))

    def foo(a, b):
        def bar(a: int, *, b: int = b):
            return a + b

        return bar(a + 1)

    assert_close(foo(*args), jit(foo)(*args))
//  [python]
def test_unpack_sequence(jit):
    def foo(tup):
        a, b = tup
        return a + b

    def bar(tup):
        a, b = map(lambda x: x, tup)  # unpack iterable
        return a + b

    jfoo = jit(foo)
    jbar = jit(bar)

    args = (4, 3)

    thunder_result = jfoo(args)
    python_result = foo(args)

    t.assert_close(thunder_result, python_result)

    thunder_result = jbar(args)
    python_result = bar(args)

    t.assert_close(thunder_result, python_result)
//  [python]
def test_unpack_sequence(jit):
    def foo(tup):
        a, b = tup
        return a + b

    def bar(tup):
        a, b = map(lambda x: x, tup)  # unpack iterable
        return a + b

    jfoo = jit(foo)
    jbar = jit(bar)

    args = (4, 3)

    thunder_result = jfoo(args)
    python_result = foo(args)

    assert_close(thunder_result, python_result)

    thunder_result = jbar(args)
    python_result = bar(args)

    assert_close(thunder_result, python_result)
//  [python]
def test_autograd_function(jit):
    class DoubleFn(torch.autograd.Function):
        @staticmethod
        def forward(ctx, a):
            return a * 2

        @staticmethod
        def backward(ctx, grad_out):
            return 2 * grad_out

    a = torch.randn(5, 5, requires_grad=True)

    def fn():
        b = DoubleFn.apply(a)
        return b

    t.assert_close(fn(), jit(fn)())
//  [python]
def test_autograd_function(jit):
    class DoubleFn(torch.autograd.Function):
        @staticmethod
        def forward(ctx, a):
            return a * 2

        @staticmethod
        def backward(ctx, grad_out):
            return 2 * grad_out

    a = torch.randn(5, 5, requires_grad=True)

    def fn():
        b = DoubleFn.apply(a)
        return b

    assert_close(fn(), jit(fn)())
//  [python]
def test_nanogpt_mlp(jit):
    from thunder.benchmarks import NanoGPTMLPBenchmark, NanoGPTConfig, _nanogpt_configs

    config: NanoGPTConfig = NanoGPTConfig(dropout=0)
    config.update(**_nanogpt_configs["gpt2"])
    bench = NanoGPTMLPBenchmark(config=config, device="cpu")
    fn = bench.fn()

    args, kwargs = bench.make_batch()

    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    t.assert_close(result, fn(*args, **kwargs))
//  [python]
def test_nanogpt_mlp(jit):
    from thunder.benchmarks import NanoGPTMLPBenchmark, NanoGPTConfig, _nanogpt_configs

    config: NanoGPTConfig = NanoGPTConfig(dropout=0)
    config.update(**_nanogpt_configs["gpt2"])
    bench = NanoGPTMLPBenchmark(config=config, device="cpu")
    fn = bench.fn()

    args, kwargs = bench.make_batch()

    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    assert_close(result, fn(*args, **kwargs))
//  [python]
def test_import(jit):
    def foo(a, b):
        import operator

        return operator.add(a, b)

    jfoo = jit(foo)

    assert jfoo(-1, 3) == foo(-1, 3)

    def foo(a, b):
        from operator import add

        return add(a, b)

    jfoo = jit(foo)

    assert jfoo(2, 7) == foo(2, 7)

    def foo(a):
        import torch.nn as nn
        from torch.nn.functional import relu

        return relu(a)

    jfoo = jit(foo)

    a = torch.randn((2, 2))

    t.assert_close(jfoo(a), foo(a))

    def foo():
        # test relative import
        from .litgpt_model import Config

        return Config

    assert jit(foo)() is foo()

    def foo():
        # test relative import
        from . import litgpt_model

        return litgpt_model.Config

    assert jit(foo)() is foo()

    # reload is implemented using exec of the module
    from . import litgpt_model
    import importlib

    importlib.reload(litgpt_model)
    assert hasattr(litgpt_model, "GPT")
//  [python]
def test_import(jit):
    def foo(a, b):
        import operator

        return operator.add(a, b)

    jfoo = jit(foo)

    assert jfoo(-1, 3) == foo(-1, 3)

    def foo(a, b):
        from operator import add

        return add(a, b)

    jfoo = jit(foo)

    assert jfoo(2, 7) == foo(2, 7)

    def foo(a):
        import torch.nn as nn
        from torch.nn.functional import relu

        return relu(a)

    jfoo = jit(foo)

    a = torch.randn((2, 2))

    assert_close(jfoo(a), foo(a))

    def foo():
        # test relative import
        from .litgpt_model import Config

        return Config

    assert jit(foo)() is foo()

    def foo():
        # test relative import
        from . import litgpt_model

        return litgpt_model.Config

    assert jit(foo)() is foo()

    # reload is implemented using exec of the module
    from . import litgpt_model
    import importlib

    importlib.reload(litgpt_model)
    assert hasattr(litgpt_model, "GPT")
//  [python]
def test_nanogpt_csa(jit):
    from thunder.benchmarks import NanoGPTCSABenchmark, NanoGPTConfig, _nanogpt_configs
    import torch as t

    config: NanoGPTConfig = NanoGPTConfig(dropout=0)
    config.update(**_nanogpt_configs["gpt2"])
    bench = NanoGPTCSABenchmark(config=config, device="cpu")
    fn = bench.fn()

    args, kwargs = bench.make_batch()

    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    t.testing.assert_close(result, fn(*args, **kwargs))
//  [python]
def test_nanogpt_csa(jit):
    from thunder.benchmarks import NanoGPTCSABenchmark, NanoGPTConfig, _nanogpt_configs

    config: NanoGPTConfig = NanoGPTConfig(dropout=0)
    config.update(**_nanogpt_configs["gpt2"])
    bench = NanoGPTCSABenchmark(config=config, device="cpu")
    fn = bench.fn()

    args, kwargs = bench.make_batch()

    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    assert_close(result, fn(*args, **kwargs))
//  [python]
def test_torch_autocast_nograd(jit):
    def fn(a, b):
        with torch.autocast("cpu"):
            return a @ b

    def fn2(a, b):
        with torch.no_grad():
            return a @ b

    a = torch.randn(5, 5)
    b = torch.randn(5, 5)
    expected = fn(a, b)
    actual = jit(fn)(a, b)

    t.assert_close(actual, expected)  # also check dtype (should be bf16)

    a = torch.randn(5, 5, requires_grad=True)
    b = torch.randn(5, 5, requires_grad=True)
    expected = fn2(a, b)
    actual = jit(fn2)(a, b)

    t.assert_close(actual, expected)
//  [python]
def test_torch_autocast_nograd(jit):
    def fn(a, b):
        with torch.autocast("cpu"):
            return a @ b

    def fn2(a, b):
        with torch.no_grad():
            return a @ b

    a = torch.randn(5, 5)
    b = torch.randn(5, 5)
    expected = fn(a, b)
    actual = jit(fn)(a, b)

    assert_close(actual, expected)  # also check dtype (should be bf16)

    a = torch.randn(5, 5, requires_grad=True)
    b = torch.randn(5, 5, requires_grad=True)
    expected = fn2(a, b)
    actual = jit(fn2)(a, b)

    assert_close(actual, expected)
//  [python]
def test_nanogpt_block(jit):
    from thunder.benchmarks import NanoGPTBlockBenchmark, NanoGPTConfig, _nanogpt_configs
    import torch as t

    config: NanoGPTConfig = NanoGPTConfig(dropout=0)
    config.update(**_nanogpt_configs["gpt2"])
    bench = NanoGPTBlockBenchmark(config=config, device="cpu")
    fn = bench.fn()

    args, kwargs = bench.make_batch()

    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    t.testing.assert_close(result, fn(*args, **kwargs))
//  [python]
def test_nanogpt_block(jit):
    from thunder.benchmarks import NanoGPTBlockBenchmark, NanoGPTConfig, _nanogpt_configs

    config: NanoGPTConfig = NanoGPTConfig(dropout=0)
    config.update(**_nanogpt_configs["gpt2"])
    bench = NanoGPTBlockBenchmark(config=config, device="cpu")
    fn = bench.fn()

    args, kwargs = bench.make_batch()

    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    assert_close(result, fn(*args, **kwargs))
//  [python]
def test_nanogpt(jit):
    from thunder.benchmarks import NanoGPTBenchmark, NanoGPTConfig, _nanogpt_configs
    import torch as t

    config: NanoGPTConfig = NanoGPTConfig(dropout=0)
    config.update(**_nanogpt_configs["test"])
    bench = NanoGPTBenchmark(config=config, device="cpu")
    fn = bench.fn()

    args, kwargs = bench.make_batch()
    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    t.testing.assert_close(result, fn(*args, **kwargs))
//  [python]
def test_nanogpt(jit):
    from thunder.benchmarks import NanoGPTBenchmark, NanoGPTConfig, _nanogpt_configs

    config: NanoGPTConfig = NanoGPTConfig(dropout=0)
    config.update(**_nanogpt_configs["test"])
    bench = NanoGPTBenchmark(config=config, device="cpu")
    fn = bench.fn()

    args, kwargs = bench.make_batch()
    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    assert_close(result, fn(*args, **kwargs))
//  [python]
def test_litgpt(jit):
    from thunder.benchmarks import LitGPTBenchmark
    from thunder.tests.litgpt_model import Config
    import torch as t

    cfg: Config = Config.from_name("gpt-neox-like")
    bench = LitGPTBenchmark(config=cfg, device="cpu", dtype=torch.bfloat16, requires_grad=True)
    fn = bench.fn()

    args, kwargs = bench.make_batch()

    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    t.testing.assert_close(result, fn(*args, **kwargs))
//  [python]
def test_litgpt(jit):
    from thunder.benchmarks import LitGPTBenchmark
    from thunder.tests.litgpt_model import Config

    cfg: Config = Config.from_name("gpt-neox-like")
    bench = LitGPTBenchmark(config=cfg, device="cpu", dtype=torch.bfloat16, requires_grad=True)
    fn = bench.fn()

    args, kwargs = bench.make_batch()

    jfn = jit(fn)
    result = jfn(*args, **kwargs)

    assert_close(result, fn(*args, **kwargs))
//  [python]
def assert_close(
    a: Tensor,
    b: Union[Tensor, List[Any]],
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
) -> None:
    """Assert that ``a`` and ``b`` are element-wise equal within a tolerance."""
    if not isinstance(b, Tensor):
        b = torch.tensor(b, device=device, dtype=a.dtype)

    torch.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
//  [python]
def assert_close(
    a: Tensor,
    b: Union[Tensor, List[Any]],
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
) -> None:
    """Assert that ``a`` and ``b`` are element-wise equal within a tolerance."""
    if not isinstance(b, Tensor):
        b = torch.tensor(b, device=device, dtype=a.dtype)

    torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
//  [python]
def assert_equal(a: Tensor, b: Union[Tensor, List[Any]]) -> None:
    """Assert that ``a`` and ``b`` are element-wise equal."""
    if not isinstance(b, Tensor):
        b = torch.tensor(b, device=device, dtype=a.dtype)

    torch.testing.assert_allclose(a, b, rtol=0, atol=0)
//  [python]
def assert_equal(a: Tensor, b: Union[Tensor, List[Any]]) -> None:
    """Assert that ``a`` and ``b`` are element-wise equal."""
    if not isinstance(b, Tensor):
        b = torch.tensor(b, device=device, dtype=a.dtype)

    torch.testing.assert_close(a, b, rtol=0, atol=0)
//  [python]
def test_binary_operations(jit):
    ops = ["+", "&", "//", "<<", "@", "*", "%", "|", "**", ">>", "-", "/", "^"]

    # NOTE Not all of the number ops support floats (for example lshift)
    number_inps = (
        (5, 9),
        (2, 8),
        (8, 2),
    )

    bool_inps = (
        (True, True),
        (False, True),
    )

    tensor_a = torch.randint(1, 10, (2, 2))
    tensor_b = torch.randint(1, 10, (2, 2))

    for op in ops:
        foo = eval(f"lambda a, b: a {op} b")
        jfoo = jit(foo)
        d = {}
        exec(f"def bar(a, b):\n a {op}= b\n return a", d)
        bar = d["bar"]
        jbar = jit(bar)

        if op != "@":
            for a, b in number_inps:
                assert jfoo(a, b) == foo(a, b)
                assert jbar(a, b) == bar(a, b)

        if op in {"&", "|", "^"}:
            for a, b in bool_inps:
                assert jfoo(a, b) == foo(a, b)
                assert jbar(a, b) == bar(a, b)

        a1 = tensor_a.clone()
        b1 = tensor_b.clone()
        expected = foo(a1, b1)
        a2 = tensor_a.clone()
        b2 = tensor_b.clone()
        actual = jfoo(a2, b2)
        t.assert_close(a1, a2)
        t.assert_close(b1, b2)
        t.assert_close(actual, expected)

        a1 = tensor_a.clone()
        b1 = tensor_b.clone()
        a2 = tensor_a.clone()
        b2 = tensor_b.clone()
        if op == "/":
            a1 = a1.float()
            a2 = a2.float()
        actual = jbar(a2, b2)
        expected = bar(a1, b1)
        t.assert_close(a1, a2)
        t.assert_close(b1, b2)
        t.assert_close(actual, expected)
//  [python]
def test_binary_operations(jit):
    ops = ["+", "&", "//", "<<", "@", "*", "%", "|", "**", ">>", "-", "/", "^"]

    # NOTE Not all of the number ops support floats (for example lshift)
    number_inps = (
        (5, 9),
        (2, 8),
        (8, 2),
    )

    bool_inps = (
        (True, True),
        (False, True),
    )

    tensor_a = torch.randint(1, 10, (2, 2))
    tensor_b = torch.randint(1, 10, (2, 2))

    for op in ops:
        foo = eval(f"lambda a, b: a {op} b")
        jfoo = jit(foo)
        d = {}
        exec(f"def bar(a, b):\n a {op}= b\n return a", d)
        bar = d["bar"]
        jbar = jit(bar)

        if op != "@":
            for a, b in number_inps:
                assert jfoo(a, b) == foo(a, b)
                assert jbar(a, b) == bar(a, b)

        if op in {"&", "|", "^"}:
            for a, b in bool_inps:
                assert jfoo(a, b) == foo(a, b)
                assert jbar(a, b) == bar(a, b)

        a1 = tensor_a.clone()
        b1 = tensor_b.clone()
        expected = foo(a1, b1)
        a2 = tensor_a.clone()
        b2 = tensor_b.clone()
        actual = jfoo(a2, b2)
        assert_close(a1, a2)
        assert_close(b1, b2)
        assert_close(actual, expected)

        a1 = tensor_a.clone()
        b1 = tensor_b.clone()
        a2 = tensor_a.clone()
        b2 = tensor_b.clone()
        if op == "/":
            a1 = a1.float()
            a2 = a2.float()
        actual = jbar(a2, b2)
        expected = bar(a1, b1)
        assert_close(a1, a2)
        assert_close(b1, b2)
        assert_close(actual, expected)
//  [python]
def test_lora_linear_works() -> None:
    lora_config = LoRAConfig(r=4, alpha=1.0, dropout_p=0.0, keys=[])

    linear_layer = Linear(8, 8, bias=True)

    lora_linear = LoRALinear(
        linear_layer, lora_config, skip_init=False, device=torch.device("cpu")
    )

    torch.nn.init.kaiming_uniform_(lora_linear.lora_B, a=math.sqrt(5))

    assert lora_linear.lora_A.shape == (4, 8)

    assert lora_linear.lora_B.shape == (8, 4)

    seqs = torch.randn([2, 8])

    orig_out = linear_layer(seqs)

    lora_out = lora_linear(seqs)

    lora_partial_out = (
        linear(seqs, lora_linear.lora_B @ lora_linear.lora_A) * lora_linear.scaling
    )

    torch.testing.assert_allclose(lora_out - orig_out, lora_partial_out)
//  [python]
def test_lora_linear_works() -> None:
    lora_config = LoRAConfig(r=4, alpha=1.0, dropout_p=0.0, keys=[])

    linear_layer = Linear(8, 8, bias=True)

    lora_linear = LoRALinear(
        linear_layer, lora_config, skip_init=False, device=torch.device("cpu")
    )

    torch.nn.init.kaiming_uniform_(lora_linear.lora_B, a=math.sqrt(5))

    assert lora_linear.lora_A.shape == (4, 8)

    assert lora_linear.lora_B.shape == (8, 4)

    seqs = torch.randn([2, 8])

    orig_out = linear_layer(seqs)

    lora_out = lora_linear(seqs)

    lora_partial_out = (
        linear(seqs, lora_linear.lora_B @ lora_linear.lora_A) * lora_linear.scaling
    )

    torch.testing.assert_close(lora_out - orig_out, lora_partial_out)
//  [python]
def test_lora_embedding_works() -> None:
    lora_config = LoRAConfig(r=4, alpha=1.0, dropout_p=0.0, keys=[])

    pad_idx = 0

    embed_layer = StandardEmbedding(4, 8, pad_idx)

    lora_embed = LoRAEmbedding(embed_layer, lora_config, device=torch.device("cpu"))

    torch.nn.init.normal_(lora_embed.lora_A)

    assert lora_embed.lora_A.shape == (4, 4)

    assert lora_embed.lora_B.shape == (8, 4)

    seqs = torch.randint(0, 4, [2, 5])

    orig_out = embed_layer(seqs)

    lora_out = lora_embed(seqs)

    lora_partial_out = embedding(
        seqs, (lora_embed.lora_B @ lora_embed.lora_A).T * lora_embed.scaling, pad_idx
    )

    torch.testing.assert_allclose(lora_out - orig_out, lora_partial_out)
//  [python]
def test_lora_embedding_works() -> None:
    lora_config = LoRAConfig(r=4, alpha=1.0, dropout_p=0.0, keys=[])

    pad_idx = 0

    embed_layer = StandardEmbedding(4, 8, pad_idx)

    lora_embed = LoRAEmbedding(embed_layer, lora_config, device=torch.device("cpu"))

    torch.nn.init.normal_(lora_embed.lora_A)

    assert lora_embed.lora_A.shape == (4, 4)

    assert lora_embed.lora_B.shape == (8, 4)

    seqs = torch.randint(0, 4, [2, 5])

    orig_out = embed_layer(seqs)

    lora_out = lora_embed(seqs)

    lora_partial_out = embedding(
        seqs, (lora_embed.lora_B @ lora_embed.lora_A).T * lora_embed.scaling, pad_idx
    )

    torch.testing.assert_close(lora_out - orig_out, lora_partial_out)
//  [python]
def test_lora_liner_merge_unmerge_work() -> None:
    lora_config = LoRAConfig(r=4, alpha=1.0, dropout_p=0.0, keys=[])

    linear_layer = Linear(8, 8, bias=True)

    lora_linear = LoRALinear(
        linear_layer, lora_config, skip_init=False, device=torch.device("cpu")
    )

    torch.nn.init.kaiming_uniform_(lora_linear.lora_B, a=math.sqrt(5))

    seqs = torch.randn([2, 8])

    orig_weight = lora_linear.weight.data.clone()

    wrapped_org_weight = lora_linear.wrapped.weight.data.clone()  # type: ignore[operator]

    torch.testing.assert_allclose(orig_weight, wrapped_org_weight)

    lora_linear.merge()

    assert lora_linear.merged

    merged_weight = lora_linear.weight.data.clone()

    wrapped_merged_weight = lora_linear.wrapped.weight.data.clone()  # type: ignore[operator]

    torch.testing.assert_allclose(merged_weight, wrapped_merged_weight)

    merged_out = lora_linear(seqs)

    lora_linear.unmerge()

    assert not lora_linear.merged

    un_merged_weight = lora_linear.weight.data.clone()

    wrapped_un_merged_weight = lora_linear.wrapped.weight.data.clone()

    torch.testing.assert_allclose(un_merged_weight, wrapped_un_merged_weight)

    un_merged_out = lora_linear(seqs)

    torch.testing.assert_allclose(orig_weight, un_merged_weight)

    torch.testing.assert_allclose(un_merged_out, merged_out)

    lora_AB = (lora_linear.lora_B @ lora_linear.lora_A * lora_linear.scaling).data

    torch.testing.assert_allclose(merged_weight - un_merged_weight, lora_AB)
//  [python]
def test_lora_liner_merge_unmerge_work() -> None:
    lora_config = LoRAConfig(r=4, alpha=1.0, dropout_p=0.0, keys=[])

    linear_layer = Linear(8, 8, bias=True)

    lora_linear = LoRALinear(
        linear_layer, lora_config, skip_init=False, device=torch.device("cpu")
    )

    torch.nn.init.kaiming_uniform_(lora_linear.lora_B, a=math.sqrt(5))

    seqs = torch.randn([2, 8])

    orig_weight = lora_linear.weight.data.clone()

    wrapped_org_weight = lora_linear.wrapped.weight.data.clone()  # type: ignore[operator]

    torch.testing.assert_close(orig_weight, wrapped_org_weight)

    lora_linear.merge()

    assert lora_linear.merged

    merged_weight = lora_linear.weight.data.clone()

    wrapped_merged_weight = lora_linear.wrapped.weight.data.clone()  # type: ignore[operator]

    torch.testing.assert_close(merged_weight, wrapped_merged_weight)

    merged_out = lora_linear(seqs)

    lora_linear.unmerge()

    assert not lora_linear.merged

    un_merged_weight = lora_linear.weight.data.clone()

    wrapped_un_merged_weight = lora_linear.wrapped.weight.data.clone()

    torch.testing.assert_close(un_merged_weight, wrapped_un_merged_weight)

    un_merged_out = lora_linear(seqs)

    torch.testing.assert_close(orig_weight, un_merged_weight)

    torch.testing.assert_close(un_merged_out, merged_out)

    lora_AB = (lora_linear.lora_B @ lora_linear.lora_A * lora_linear.scaling).data

    torch.testing.assert_close(merged_weight - un_merged_weight, lora_AB)
//  [python]
def test_add_mlp_layer(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        max_hidden_layers=5,
        device=device,
    )

    initial_hidden_size = len(evolvable_cnn.hidden_size)
    initial_net = evolvable_cnn.value_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.add_mlp_layer()
    new_net = evolvable_cnn.value_net
    if initial_hidden_size < 10:
        assert len(evolvable_cnn.hidden_size) == initial_hidden_size + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.allclose(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.hidden_size) == initial_hidden_size
//  [python]
def test_add_mlp_layer(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        max_hidden_layers=5,
        device=device,
    )

    initial_hidden_size = len(evolvable_cnn.hidden_size)
    initial_net = evolvable_cnn.value_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.add_mlp_layer()
    new_net = evolvable_cnn.value_net
    if initial_hidden_size < 10:
        assert len(evolvable_cnn.hidden_size) == initial_hidden_size + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.hidden_size) == initial_hidden_size
//  [python]
def test_lora_embedding_merge_unmerge_work() -> None:
    lora_config = LoRAConfig(r=4, alpha=1.0, dropout_p=0.0, keys=[])

    pad_idx = 0

    embed_layer = StandardEmbedding(4, 8, pad_idx)

    lora_embed = LoRAEmbedding(embed_layer, lora_config, device=torch.device("cpu"))

    torch.nn.init.normal_(lora_embed.lora_A)

    seqs = torch.randint(0, 4, [2, 5])

    orig_weight = lora_embed.weight.data.clone()

    wrapped_orig_weight = lora_embed.wrapped.weight.data.clone()  # type: ignore[operator]

    torch.testing.assert_allclose(orig_weight, wrapped_orig_weight)

    lora_embed.merge()

    assert lora_embed.merged

    merged_weight = lora_embed.weight.data.clone()

    wrapped_merged_weight = lora_embed.wrapped.weight.data.clone()  # type: ignore[operator]

    torch.testing.assert_allclose(merged_weight, wrapped_merged_weight)

    merged_out = lora_embed(seqs)

    lora_embed.unmerge()

    assert not lora_embed.merged

    un_merged_weight = lora_embed.weight.data.clone()

    wrapped_un_merged_weight = lora_embed.wrapped.weight.data.clone()

    torch.testing.assert_allclose(un_merged_weight, wrapped_un_merged_weight)

    un_merged_out = lora_embed(seqs)

    torch.testing.assert_allclose(orig_weight, un_merged_weight)

    torch.testing.assert_allclose(un_merged_out, merged_out)

    lora_AB = (lora_embed.lora_B @ lora_embed.lora_A * lora_embed.scaling).data.T

    torch.testing.assert_allclose(merged_weight - un_merged_weight, lora_AB)
//  [python]
def test_lora_embedding_merge_unmerge_work() -> None:
    lora_config = LoRAConfig(r=4, alpha=1.0, dropout_p=0.0, keys=[])

    pad_idx = 0

    embed_layer = StandardEmbedding(4, 8, pad_idx)

    lora_embed = LoRAEmbedding(embed_layer, lora_config, device=torch.device("cpu"))

    torch.nn.init.normal_(lora_embed.lora_A)

    seqs = torch.randint(0, 4, [2, 5])

    orig_weight = lora_embed.weight.data.clone()

    wrapped_orig_weight = lora_embed.wrapped.weight.data.clone()  # type: ignore[operator]

    torch.testing.assert_close(orig_weight, wrapped_orig_weight)

    lora_embed.merge()

    assert lora_embed.merged

    merged_weight = lora_embed.weight.data.clone()

    wrapped_merged_weight = lora_embed.wrapped.weight.data.clone()  # type: ignore[operator]

    torch.testing.assert_close(merged_weight, wrapped_merged_weight)

    merged_out = lora_embed(seqs)

    lora_embed.unmerge()

    assert not lora_embed.merged

    un_merged_weight = lora_embed.weight.data.clone()

    wrapped_un_merged_weight = lora_embed.wrapped.weight.data.clone()

    torch.testing.assert_close(un_merged_weight, wrapped_un_merged_weight)

    un_merged_out = lora_embed(seqs)

    torch.testing.assert_close(orig_weight, un_merged_weight)

    torch.testing.assert_close(un_merged_out, merged_out)

    lora_AB = (lora_embed.lora_B @ lora_embed.lora_A * lora_embed.scaling).data.T

    torch.testing.assert_close(merged_weight - un_merged_weight, lora_AB)
//  [python]
def test_remove_mlp_layer(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        max_hidden_layers=5,
        min_hidden_layers=1,
        device=device,
    )

    initial_hidden_size = len(evolvable_cnn.hidden_size)
    initial_net = evolvable_cnn.value_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.remove_mlp_layer()
    new_net = evolvable_cnn.value_net
    if initial_hidden_size > 1:
        assert len(evolvable_cnn.hidden_size) == initial_hidden_size - 1
        for key, param in new_net.named_parameters():
            if (
                key in initial_net_dict.keys()
                and param.shape == initial_net_dict[key].shape
            ):
                torch.allclose(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.hidden_size) == initial_hidden_size
//  [python]
def test_remove_mlp_layer(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        max_hidden_layers=5,
        min_hidden_layers=1,
        device=device,
    )

    initial_hidden_size = len(evolvable_cnn.hidden_size)
    initial_net = evolvable_cnn.value_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.remove_mlp_layer()
    new_net = evolvable_cnn.value_net
    if initial_hidden_size > 1:
        assert len(evolvable_cnn.hidden_size) == initial_hidden_size - 1
        for key, param in new_net.named_parameters():
            if (
                key in initial_net_dict.keys()
                and param.shape == initial_net_dict[key].shape
            ):
                torch.testing.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.hidden_size) == initial_hidden_size
//  [python]
def test_add_cnn_layer(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        device=device,
    )
    initial_channel_num = len(evolvable_cnn.channel_size)
    initial_net = evolvable_cnn.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.add_cnn_layer()
    new_net = evolvable_cnn.feature_net
    if initial_channel_num < 6:
        assert len(evolvable_cnn.channel_size) == initial_channel_num + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.allclose(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.channel_size) == initial_channel_num
//  [python]
def test_add_cnn_layer(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        device=device,
    )
    initial_channel_num = len(evolvable_cnn.channel_size)
    initial_net = evolvable_cnn.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.add_cnn_layer()
    new_net = evolvable_cnn.feature_net
    if initial_channel_num < 6:
        assert len(evolvable_cnn.channel_size) == initial_channel_num + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.channel_size) == initial_channel_num
//  [python]
def test_prepare_prompt():
    model_runner = ModelRunner(None, None, None)
    model_runner.set_block_size(16)

    batch_size = random.randint(1, 256)
    prompt_lens = []
    seq_group_metadata_list = []
    for i in range(batch_size):
        # make sure all tokens fit into one block
        prompt_len = i % (model_runner.block_size - 1) + 1
        prompt_lens.append(prompt_len)
        seq_data = list(range(prompt_len))
        seq_group_metadata_list.append(
            SequenceGroupMetadata(
                request_id=f"test_{i}",
                is_prompt=True,
                seq_data={0: SequenceData(seq_data)},
                sampling_params=SamplingParams(temperature=0),
                block_tables={0: [1]},
            ))

    expected_selected_token_indices = []
    selected_token_start_idx = 0
    max_seq_len = max(prompt_lens)
    for prompt_len in prompt_lens:
        expected_selected_token_indices.append(selected_token_start_idx +
                                               prompt_len - 1)
        selected_token_start_idx += max_seq_len
    input_tokens, input_positions, _ = model_runner._prepare_prompt(
        seq_group_metadata_list)
    sampling_metadata = model_runner._prepare_sample(seq_group_metadata_list,
                                                     prompt_lens)
    assert input_tokens.shape == (batch_size, max_seq_len)
    assert input_positions.shape == (batch_size, max_seq_len)
    t.assert_close(input_tokens, input_positions)

    actual = sampling_metadata.selected_token_indices
    expected = torch.tensor(expected_selected_token_indices,
                            device=actual.device,
                            dtype=actual.dtype)
    t.assert_close(actual, expected)