jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_lookaside_bool():
    def foo(a, b, i):
        if bool(i):
            return a + b
        return a - b

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    expected = foo(a, b, 0)
    actual = jfoo(a, b, 0)
    assert_close(expected, actual)

    expected = foo(a, b, 1)
    actual = jfoo(a, b, 1)
    assert_close(expected, actual)
//  [python]
def test_nn_module():
    m = torch.nn.Linear(3, 4)
    m2 = torch.nn.Sequential(
        torch.nn.Linear(3, 4),
        torch.nn.Linear(4, 3),
    )
    a = torch.randn(2, 3)

    tom = thunder.jit(m)
    assert isinstance(tom, thunder.ThunderModule)
    # attributes are forwarded
    assert isinstance(tom.weight, torch.Tensor)
    expected = m(a)
    actual = tom(a)
    assert_close(expected, actual)

    tom2 = thunder.jit(m2)
    assert isinstance(tom2, thunder.ThunderModule)
    # `ThunderModule` is not subscriptable even though it compiles a `Sequential`
    with pytest.raises(TypeError, match="not subscriptable"):
        assert isinstance(tom2[1].weight, torch.Tensor)
    # saving is the same
    torch.testing.assert_allclose(tom2.state_dict(), m2.state_dict(), rtol=0, atol=0)
    # loading works
    tom2.load_state_dict(m2.state_dict(), strict=True, assign=True)
    expected = m2(a)
    actual = tom2(a)
    assert_close(expected, actual)

    def fn(a):
        return m(a)

    jfn = thunder.jit(fn)
    expected = fn(a)
    actual = jfn(a)
    assert_close(expected, actual)

    def fn2(a):
        return m2(a)

    jfn2 = thunder.jit(fn2)
    expected = fn2(a)
    actual = jfn2(a)
    assert_close(expected, actual)
//  [python]
def test_nn_module():
    m = torch.nn.Linear(3, 4)
    m2 = torch.nn.Sequential(
        torch.nn.Linear(3, 4),
        torch.nn.Linear(4, 3),
    )
    a = torch.randn(2, 3)

    tom = thunder.jit(m)
    assert isinstance(tom, thunder.ThunderModule)
    # attributes are forwarded
    assert isinstance(tom.weight, torch.Tensor)
    expected = m(a)
    actual = tom(a)
    assert_close(expected, actual)

    tom2 = thunder.jit(m2)
    assert isinstance(tom2, thunder.ThunderModule)
    # `ThunderModule` is not subscriptable even though it compiles a `Sequential`
    with pytest.raises(TypeError, match="not subscriptable"):
        assert isinstance(tom2[1].weight, torch.Tensor)
    # saving is the same
    torch.testing.assert_close(tom2.state_dict(), m2.state_dict(), rtol=0, atol=0)
    # loading works
    tom2.load_state_dict(m2.state_dict(), strict=True, assign=True)
    expected = m2(a)
    actual = tom2(a)
    assert_close(expected, actual)

    def fn(a):
        return m(a)

    jfn = thunder.jit(fn)
    expected = fn(a)
    actual = jfn(a)
    assert_close(expected, actual)

    def fn2(a):
        return m2(a)

    jfn2 = thunder.jit(fn2)
    expected = fn2(a)
    actual = jfn2(a)
    assert_close(expected, actual)
//  [python]
def test_nanogpt_block():
    from thunder.benchmarks import NanoGPTBlockBenchmark, NanoGPTConfig, _nanogpt_configs

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

    args, kwargs = bench.make_batch()

    jfn = thunder.jit(module)
    result = jfn(*args, **kwargs)

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

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

    args, kwargs = bench.make_batch()

    jfn = thunder.jit(module)
    result = jfn(*args, **kwargs)

    assert_close(result, module(*args, **kwargs))
//  [python]
def test_proxy_no_multiple_renames(device):
    if device == "cuda" and not torch.cuda.is_available():
        pytest.skip("CUDA not available")

    device = torch.device(device)

    def f(a, b, c, d):
        def g(b, a):
            return b + a

        def h(d, c):
            return d + c

        return g(a, b) + h(d, c)

    jf = thunder.jit(f)

    a = torch.rand(1, device=device)
    b = torch.rand(1, device=device)
    c = torch.rand(1, device=device)
    d = torch.rand(1, device=device)

    expected = f(a, b, c, d)
    actual = jf(a, b, c, d)
    t.assert_close(expected, actual)

    comp_trace = thunder.last_traces(jf)[-1]
    args_names = tuple(p.name for p in comp_trace.args)
    # Once proxies got their names in `f`, these should not change in `g` and `h`
    assert args_names == ("a", "b", "c", "d")
//  [python]
def test_proxy_no_multiple_renames(device):
    if device == "cuda" and not torch.cuda.is_available():
        pytest.skip("CUDA not available")

    device = torch.device(device)

    def f(a, b, c, d):
        def g(b, a):
            return b + a

        def h(d, c):
            return d + c

        return g(a, b) + h(d, c)

    jf = thunder.jit(f)

    a = torch.rand(1, device=device)
    b = torch.rand(1, device=device)
    c = torch.rand(1, device=device)
    d = torch.rand(1, device=device)

    expected = f(a, b, c, d)
    actual = jf(a, b, c, d)
    assert_close(expected, actual)

    comp_trace = thunder.last_traces(jf)[-1]
    args_names = tuple(p.name for p in comp_trace.args)
    # Once proxies got their names in `f`, these should not change in `g` and `h`
    assert args_names == ("a", "b", "c", "d")
//  [python]
def test_cache_basic():
    def foo(a, b):
        return a + b

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    expected = foo(a, b)
    actual = jfoo(a, b)
    torch.testing.assert_allclose(expected, actual)
    assert thunder.cache_misses(jfoo) == 1
    assert thunder.cache_hits(jfoo) == 0

    expected = foo(a, b)
    actual = jfoo(a, b)
    torch.testing.assert_allclose(expected, actual)
    assert thunder.cache_misses(jfoo) == 1
    assert thunder.cache_hits(jfoo) == 1

    # Tests rank changing
    a = torch.randn((2), device="cpu")

    expected = foo(a, b)
    actual = jfoo(a, b)
    torch.testing.assert_allclose(expected, actual)
    assert thunder.cache_misses(jfoo) == 2
    assert thunder.cache_hits(jfoo) == 1

    expected = foo(a, b)
    actual = jfoo(a, b)
    torch.testing.assert_allclose(expected, actual)
    assert thunder.cache_misses(jfoo) == 2
    assert thunder.cache_hits(jfoo) == 2

    # Tests dtype changing
    a = torch.randn((2, 2), device="cpu", dtype=torch.bfloat16)
    b = torch.randn((2, 2), device="cpu", dtype=torch.bfloat16)

    expected = foo(a, b)
    actual = jfoo(a, b)
    torch.testing.assert_allclose(expected, actual)
    assert thunder.cache_misses(jfoo) == 3
    assert thunder.cache_hits(jfoo) == 2

    expected = foo(a, b)
    actual = jfoo(a, b)
    torch.testing.assert_allclose(expected, actual)
    assert thunder.cache_misses(jfoo) == 3
    assert thunder.cache_hits(jfoo) == 3

    # Tests shape changing
    a = torch.randn((2, 1), device="cpu", dtype=torch.bfloat16)
    b = torch.randn((2, 1), device="cpu", dtype=torch.bfloat16)

    expected = foo(a, b)
    actual = jfoo(a, b)
    torch.testing.assert_allclose(expected, actual)
    assert thunder.cache_misses(jfoo) == 4
    assert thunder.cache_hits(jfoo) == 3

    expected = foo(a, b)
    actual = jfoo(a, b)
    torch.testing.assert_allclose(expected, actual)
    assert thunder.cache_misses(jfoo) == 4
    assert thunder.cache_hits(jfoo) == 4
//  [python]
def test_cache_basic():
    def foo(a, b):
        return a + b

    jfoo = thunder.jit(foo)

    a = torch.randn((2, 2), device="cpu")
    b = torch.randn((2, 2), device="cpu")

    expected = foo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 1
    assert thunder.cache_hits(jfoo) == 0

    expected = foo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 1
    assert thunder.cache_hits(jfoo) == 1

    # Tests rank changing
    a = torch.randn((2), device="cpu")

    expected = foo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 2
    assert thunder.cache_hits(jfoo) == 1

    expected = foo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 2
    assert thunder.cache_hits(jfoo) == 2

    # Tests dtype changing
    a = torch.randn((2, 2), device="cpu", dtype=torch.bfloat16)
    b = torch.randn((2, 2), device="cpu", dtype=torch.bfloat16)

    expected = foo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 3
    assert thunder.cache_hits(jfoo) == 2

    expected = foo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 3
    assert thunder.cache_hits(jfoo) == 3

    # Tests shape changing
    a = torch.randn((2, 1), device="cpu", dtype=torch.bfloat16)
    b = torch.randn((2, 1), device="cpu", dtype=torch.bfloat16)

    expected = foo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 4
    assert thunder.cache_hits(jfoo) == 3

    expected = foo(a, b)
    actual = jfoo(a, b)
    assert_close(expected, actual)
    assert thunder.cache_misses(jfoo) == 4
    assert thunder.cache_hits(jfoo) == 4
//  [python]
def test_litgpt():
    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)
    module = bench.fn()
    jfn = thunder.jit(module)

    args, kwargs = bench.make_batch()

    # the transforms of thunder_backward introduce numerical differences in bfloat16 mode that cause flakiness
    for p in module.parameters():
        p.requires_grad_(False)

    reference = module(*args, **kwargs)
    result = jfn(*args, **kwargs)
    t.testing.assert_close(result, reference)

    args, kwargs = bench.make_batch()
    result = jfn(*args, **kwargs)
    t.testing.assert_close(result, module(*args, **kwargs))

    assert thunder.cache_misses(jfn) == 1
    assert thunder.cache_hits(jfn) == 1
//  [python]
def test_litgpt():
    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)
    module = bench.fn()
    jfn = thunder.jit(module)

    args, kwargs = bench.make_batch()

    # the transforms of thunder_backward introduce numerical differences in bfloat16 mode that cause flakiness
    for p in module.parameters():
        p.requires_grad_(False)

    reference = module(*args, **kwargs)
    result = jfn(*args, **kwargs)
    assert_close(result, reference)

    args, kwargs = bench.make_batch()
    result = jfn(*args, **kwargs)
    assert_close(result, module(*args, **kwargs))

    assert thunder.cache_misses(jfn) == 1
    assert thunder.cache_hits(jfn) == 1
//  [python]
def test_nanogpt_mlp():
    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")
    module = bench.fn().mlp

    args, kwargs = bench.make_batch()

    jfn = thunder.jit(module)
    result = jfn(*args, **kwargs)

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

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

    args, kwargs = bench.make_batch()

    jfn = thunder.jit(module)
    result = jfn(*args, **kwargs)

    assert_close(result, module(*args, **kwargs))
//  [python]
def test_nanogpt_attn():
    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")
    module = bench.fn()
    module = module.attn

    # the transforms of thunder_backward introduce numerical differences in bfloat16 mode that cause flakiness
    for p in module.parameters():
        p.requires_grad_(False)

    args, kwargs = bench.make_batch()

    jfn = thunder.jit(module)
    result = jfn(*args, **kwargs)

    t.testing.assert_close(result, module(*args, **kwargs), atol=3e-5, rtol=1e-5)
//  [python]
def test_nanogpt_attn():
    from thunder.benchmarks import NanoGPTBlockBenchmark, NanoGPTConfig, _nanogpt_configs

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

    # the transforms of thunder_backward introduce numerical differences in bfloat16 mode that cause flakiness
    for p in module.parameters():
        p.requires_grad_(False)

    args, kwargs = bench.make_batch()

    jfn = thunder.jit(module)
    result = jfn(*args, **kwargs)

    assert_close(result, module(*args, **kwargs), atol=3e-5, rtol=1e-5)
//  [python]
def test_nanogpt():
    from thunder.benchmarks import NanoGPTBenchmark, NanoGPTConfig, _nanogpt_configs

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

    # the transforms of thunder_backward introduce numerical differences in bfloat16 mode that cause flakiness
    for p in module.parameters():
        p.requires_grad_(False)

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

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

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

    # the transforms of thunder_backward introduce numerical differences in bfloat16 mode that cause flakiness
    for p in module.parameters():
        p.requires_grad_(False)

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

    assert_close(result, module(*args, **kwargs))
//  [python]
def test_litgpt_variants(name, device):
    if device == "cuda" and not torch.cuda.is_available():
        pytest.skip("CUDA not available")

    device = torch.device(device)

    x = torch.randint(0, 200, (5, 5), device=device)
    config = litgpt_model.Config.from_name(name)

    with device:
        reference = litgpt_model.GPT(config)
    expected_logits = reference(x)

    expected_logits.sum().backward()

    with device:
        model = litgpt_model.GPT(config)
    model.load_state_dict(reference.state_dict())
    tom = thunder.jit(model, executors=nvfuserex if device.type == "cuda" else torchex)
    actual_logits = tom(x)
    t.assert_close(actual_logits, expected_logits)

    actual_logits.sum().backward()

    for param1, param2 in zip(reference.parameters(), model.parameters()):
        assert param1 is not param2
        assert param1.grad is not None
        t.assert_close(param1.grad, param2.grad, rtol=1e-2, atol=1e-2)
//  [python]
def test_litgpt_variants(name, device):
    if device == "cuda" and not torch.cuda.is_available():
        pytest.skip("CUDA not available")

    device = torch.device(device)

    x = torch.randint(0, 200, (5, 5), device=device)
    config = litgpt_model.Config.from_name(name)

    with device:
        reference = litgpt_model.GPT(config)
    expected_logits = reference(x)

    expected_logits.sum().backward()

    with device:
        model = litgpt_model.GPT(config)
    model.load_state_dict(reference.state_dict())
    tom = thunder.jit(model, executors=nvfuserex if device.type == "cuda" else torchex)
    actual_logits = tom(x)
    assert_close(actual_logits, expected_logits)

    actual_logits.sum().backward()

    for param1, param2 in zip(reference.parameters(), model.parameters()):
        assert param1 is not param2
        assert param1.grad is not None
        torch.testing.assert_close(param1.grad, param2.grad, rtol=1e-2, atol=1e-2)
//  [python]
def test_at_k(self) -> None:
        """
        Test NDCG @ K.
        """
        model_output = get_test_case_at_k()
        self.ndcg_at_k.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.ndcg_at_k.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_at_k(self) -> None:
        """
        Test NDCG @ K.
        """
        model_output = get_test_case_at_k()
        self.ndcg_at_k.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.ndcg_at_k.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_litgpt_variants_kvcache(name, device):
    import torch._dynamo  # this monkeypatches torch.manual_seed

    if device == "cuda" and not torch.cuda.is_available():
        pytest.skip("CUDA not available")

    device = torch.device(device)
    x = torch.randint(0, 200, (1, 2), device=device)
    config = litgpt_model.Config.from_name(name)

    with device:
        model = litgpt_model.GPT(config)
        model.max_seq_length = 3

    for p in model.parameters():
        p.requires_grad_(False)

    executors = nvfuserex if device.type == "cuda" else torchex
    executors = [sdpa_ex] + executors

    def sample(logits):
        return torch.argmax(logits[:, -1], dim=-1, keepdim=True)

    # the reference is 2 regular forward without the kv cache
    logits_1 = model(x)
    token_1 = sample(logits_1)
    logits_2 = model(torch.cat((x, token_1), dim=-1))

    with device:
        model.set_kv_cache(batch_size=1)
    tom = thunder.jit(model, executors=executors)  # , disable_torch_autograd_support=True

    # kv cache prefill
    thunder_logits_1 = tom(x, torch.tensor([0, 1], device=device))
    thunder_token_1 = sample(thunder_logits_1)
    # 1 token generation
    thunder_logits_2 = tom(thunder_token_1, torch.tensor([2], device=device))

    t.assert_close(logits_1, thunder_logits_1)
    t.assert_close(logits_2[:, -1:], thunder_logits_2)
//  [python]
def test_litgpt_variants_kvcache(name, device):
    import torch._dynamo  # this monkeypatches torch.manual_seed

    if device == "cuda" and not torch.cuda.is_available():
        pytest.skip("CUDA not available")

    device = torch.device(device)
    x = torch.randint(0, 200, (1, 2), device=device)
    config = litgpt_model.Config.from_name(name)

    with device:
        model = litgpt_model.GPT(config)
        model.max_seq_length = 3

    for p in model.parameters():
        p.requires_grad_(False)

    executors = nvfuserex if device.type == "cuda" else torchex
    executors = [sdpa_ex] + executors

    def sample(logits):
        return torch.argmax(logits[:, -1], dim=-1, keepdim=True)

    # the reference is 2 regular forward without the kv cache
    logits_1 = model(x)
    token_1 = sample(logits_1)
    logits_2 = model(torch.cat((x, token_1), dim=-1))

    with device:
        model.set_kv_cache(batch_size=1)
    tom = thunder.jit(model, executors=executors)  # , disable_torch_autograd_support=True

    # kv cache prefill
    thunder_logits_1 = tom(x, torch.tensor([0, 1], device=device))
    thunder_token_1 = sample(thunder_logits_1)
    # 1 token generation
    thunder_logits_2 = tom(thunder_token_1, torch.tensor([2], device=device))

    assert_close(logits_1, thunder_logits_1)
    assert_close(logits_2[:, -1:], thunder_logits_2)
//  [python]
def test_remove_single_length_sessions(self) -> None:
        """
        Test NDCG with removing single length sessions.
        """
        model_output = get_test_case_remove_single_length_sessions()
        self.ndcg_remove_single_length_sessions.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.ndcg_remove_single_length_sessions.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_remove_single_length_sessions(self) -> None:
        """
        Test NDCG with removing single length sessions.
        """
        model_output = get_test_case_remove_single_length_sessions()
        self.ndcg_remove_single_length_sessions.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.ndcg_remove_single_length_sessions.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_multiple_sessions(self) -> None:
        """
        Test multiple sessions in a single update.
        """
        model_output = get_test_case_multiple_sessions_within_batch()
        self.non_exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.non_exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )

        self.exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_multiple_sessions(self) -> None:
        """
        Test multiple sessions in a single update.
        """
        model_output = get_test_case_multiple_sessions_within_batch()
        self.non_exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.non_exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )

        self.exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_single_session(self) -> None:
        """
        Test single session in a update.
        """
        model_output = get_test_case_multiple_sessions_within_batch()
        self.non_exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.non_exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )

        self.exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_single_session(self) -> None:
        """
        Test single session in a update.
        """
        model_output = get_test_case_multiple_sessions_within_batch()
        self.non_exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.non_exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )

        self.exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_negative_sessions(self) -> None:
        """
        Test sessions where all labels are 0.
        """
        model_output = get_test_case_all_labels_zero()
        self.non_exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.non_exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )

        self.exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_negative_sessions(self) -> None:
        """
        Test sessions where all labels are 0.
        """
        model_output = get_test_case_all_labels_zero()
        self.non_exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.non_exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )

        self.exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_another_multiple_sessions(self) -> None:
        """
        Test another multiple sessions in a single update.
        """
        model_output = get_test_case_another_multiple_sessions_within_batch()
        self.non_exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.non_exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )

        self.exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_another_multiple_sessions(self) -> None:
        """
        Test another multiple sessions in a single update.
        """
        model_output = get_test_case_another_multiple_sessions_within_batch()
        self.non_exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.non_exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )

        self.exponential_ndcg.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )
        metric = self.exponential_ndcg.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_apply_negative_task_mask(self) -> None:
        """
        Test NDCG with apply negative task mask.
        """
        model_output = get_test_case_negative_task()
        self.ndcg_apply_negative_task_mask.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )

        metric = self.ndcg_apply_negative_task_mask.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_apply_negative_task_mask(self) -> None:
        """
        Test NDCG with apply negative task mask.
        """
        model_output = get_test_case_negative_task()
        self.ndcg_apply_negative_task_mask.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )

        metric = self.ndcg_apply_negative_task_mask.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_build_teacher_torchscript(self):
        """Check can build teacher using torchscript fname in config"""
        # create torchscript
        model = DivideInputBy2()
        traced_model = torch.jit.trace(model, torch.randn(5))
        with make_temp_directory("tmp") as output_dir:
            fname = f"{output_dir}/tmp.pt"
            torch.jit.save(traced_model, fname)

            # create teacher
            cfg = _get_default_cfg()
            cfg.DISTILLATION.TEACHER.TORCHSCRIPT_FNAME = fname
            teacher = _build_teacher(cfg)
            batched_inputs = torch.randn(5)
            gt = batched_inputs / 2.0
            output = teacher(batched_inputs)
            torch.testing.assert_allclose(torch.Tensor(output), gt)
//  [python]
def test_build_teacher_torchscript(self):
        """Check can build teacher using torchscript fname in config"""
        # create torchscript
        model = DivideInputBy2()
        traced_model = torch.jit.trace(model, torch.randn(5))
        with make_temp_directory("tmp") as output_dir:
            fname = f"{output_dir}/tmp.pt"
            torch.jit.save(traced_model, fname)

            # create teacher
            cfg = _get_default_cfg()
            cfg.DISTILLATION.TEACHER.TORCHSCRIPT_FNAME = fname
            teacher = _build_teacher(cfg)
            batched_inputs = torch.randn(5)
            gt = batched_inputs / 2.0
            output = teacher(batched_inputs)
            torch.testing.assert_close(torch.Tensor(output), gt)
//  [python]
def test_cached_layer_tensor(self):
        """Check cached layer saves layer output"""
        model = AddOne()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = torch.randn(1)
        output = model(input)
        torch.testing.assert_allclose(output, cache["test_layer"])
//  [python]
def test_cached_layer_tensor(self):
        """Check cached layer saves layer output"""
        model = AddOne()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = torch.randn(1)
        output = model(input)
        torch.testing.assert_close(output, cache["test_layer"])
//  [python]
def test_build_teacher_torchscript_gpu(self):
        """Check teacher moved to cuda"""
        model = AddOne()
        traced_model = torch.jit.trace(model, torch.randn(5))
        with make_temp_directory("tmp") as output_dir:
            fname = f"{output_dir}/tmp.pt"
            torch.jit.save(traced_model, fname)

            # create teacher
            cfg = _get_default_cfg()
            cfg.MODEL.DEVICE = "cuda"
            cfg.DISTILLATION.TEACHER.TORCHSCRIPT_FNAME = fname
            teacher = _build_teacher(cfg)
            batched_inputs = torch.randn(5).to("cuda")
            gt = batched_inputs + torch.Tensor([1]).to("cuda")
            output = teacher(batched_inputs)
            torch.testing.assert_allclose(torch.Tensor(output), gt)
//  [python]
def test_build_teacher_torchscript_gpu(self):
        """Check teacher moved to cuda"""
        model = AddOne()
        traced_model = torch.jit.trace(model, torch.randn(5))
        with make_temp_directory("tmp") as output_dir:
            fname = f"{output_dir}/tmp.pt"
            torch.jit.save(traced_model, fname)

            # create teacher
            cfg = _get_default_cfg()
            cfg.MODEL.DEVICE = "cuda"
            cfg.DISTILLATION.TEACHER.TORCHSCRIPT_FNAME = fname
            teacher = _build_teacher(cfg)
            batched_inputs = torch.randn(5).to("cuda")
            gt = batched_inputs + torch.Tensor([1]).to("cuda")
            output = teacher(batched_inputs)
            torch.testing.assert_close(torch.Tensor(output), gt)
//  [python]
def test_cached_layer_list(self):
        """Check cached layer saves list"""
        model = DivideInputBy2()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = [torch.randn(1) for _ in range(2)]
        output = model(input)
        t.assert_close(output, cache["test_layer"])
//  [python]
def test_cached_layer_list(self):
        """Check cached layer saves list"""
        model = DivideInputBy2()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = [torch.randn(1) for _ in range(2)]
        output = model(input)
        torch.testing.assert_close(output, cache["test_layer"])
//  [python]
def test_case_report_as_increasing_ndcg_and_scale_by_weights_tensor(self) -> None:
        """
        Test NDCG with reporting as increasing NDCG and scaling by weights tensor correctly.
        """
        model_output = get_test_case_scale_by_weights_tensor()
        self.ndcg_report_as_increasing_and_scale_by_weights_tensor.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )

        metric = self.ndcg_report_as_increasing_and_scale_by_weights_tensor.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_case_report_as_increasing_ndcg_and_scale_by_weights_tensor(self) -> None:
        """
        Test NDCG with reporting as increasing NDCG and scaling by weights tensor correctly.
        """
        model_output = get_test_case_scale_by_weights_tensor()
        self.ndcg_report_as_increasing_and_scale_by_weights_tensor.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
            required_inputs={SESSION_KEY: model_output["session_ids"][0]},
        )

        metric = self.ndcg_report_as_increasing_and_scale_by_weights_tensor.compute()
        actual_metric = metric[f"ndcg-{DefaultTaskInfo.name}|lifetime_ndcg"]
        expected_metric = model_output["expected_ndcg_non_exp"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_cached_layer_tuple(self):
        """Check cached layer saves list"""
        model = DivideInputBy2()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = (torch.randn(1) for _ in range(2))
        output = model(input)
        torch.testing.assert_allclose(output, cache["test_layer"])
//  [python]
def test_cached_layer_tuple(self):
        """Check cached layer saves list"""
        model = DivideInputBy2()
        cache = {}
        dynamic_mixin(
            model,
            CachedLayer,
            init_dict={"label": "test_layer", "cache": cache},
        )
        input = (torch.randn(1) for _ in range(2))
        output = model(input)
        torch.testing.assert_close(output, cache["test_layer"])