jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_integer_isinstance_mimicry(executor, device: str, dtype: dtypes.dtype):
    # isinstance() works as expected
    def foo(a, b, c):
        if isinstance(a, int):
            return clang.add(a, b)

        return clang.add(b, c)

    traced_foo = executor.make_callable(foo)

    tdtype = ltorch.to_torch_dtype(dtype)
    a = make_tensor((2, 1), device=device, dtype=tdtype)
    b = make_tensor((2, 2), device=device, dtype=tdtype)
    c = make_tensor((1, 2), device=device, dtype=tdtype)

    thunder_result = traced_foo(a, b, c)
    torch_result = b + c
    assert_close(thunder_result, torch_result)

    thunder_result = traced_foo(2, b, c)
    torch_result = 2 + b
    assert_close(thunder_result, torch_result)

    # type() doesn't work (it returns the actual type)
    def bar(a, b, c):
        if type(a) is int:
            return clang.add(a, b)

        return clang.add(b, c)

    traced_bar = executor.make_callable(bar)

    try:
        thunder_result = traced_bar(a, b, c)
        torch_result = b + c
        assert_close(thunder_result, torch_result)
        pytest.fail()
    except BaseException:
        pass

    try:
        thunder_result = traced_bar(2, b, c)
        torch_result = 2 + b
        assert_close(thunder_result, torch_result)
        pytest.fail()
    except BaseException:
        pass
//  [python]
def test_varargs_and_kwargs(executor, device, dtype):
    def foo(a, b, *posargs, e, **kwargs):
        accum = a
        for x in posargs:
            accum = a + x

        d = b + e + kwargs["f"]

        return accum, d, kwargs["g"]

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    b = make_tensor((2, 2, 2), device=device, dtype=tdtype)
    c = make_tensor((2, 2), device=device, dtype=tdtype)
    d = make_tensor((2,), device=device, dtype=tdtype)
    e = make_tensor((2,), device=device, dtype=tdtype)
    f = make_tensor((2,), device=device, dtype=tdtype)
    g = make_tensor((2,), device=device, dtype=tdtype)

    thunder_result = traced_foo(a, b, c, d, e=e, f=f, g=g)
    torch_result = foo(a, b, c, d, e=e, f=f, g=g)

    t.assert_close(thunder_result, torch_result)
//  [python]
def test_varargs_and_kwargs(executor, device, dtype):
    def foo(a, b, *posargs, e, **kwargs):
        accum = a
        for x in posargs:
            accum = a + x

        d = b + e + kwargs["f"]

        return accum, d, kwargs["g"]

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    b = make_tensor((2, 2, 2), device=device, dtype=tdtype)
    c = make_tensor((2, 2), device=device, dtype=tdtype)
    d = make_tensor((2,), device=device, dtype=tdtype)
    e = make_tensor((2,), device=device, dtype=tdtype)
    f = make_tensor((2,), device=device, dtype=tdtype)
    g = make_tensor((2,), device=device, dtype=tdtype)

    thunder_result = traced_foo(a, b, c, d, e=e, f=f, g=g)
    torch_result = foo(a, b, c, d, e=e, f=f, g=g)

    assert_close(thunder_result, torch_result)
//  [python]
def test_model_7b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = MistralForCausalLM.from_pretrained(
            "mistralai/Mistral-7B-v0.1", device_map="auto", torch_dtype=torch.float16
        )
        input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
        with torch.no_grad():
            out = model(input_ids).logits.cpu()
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-2.5548, -2.5737, -3.0600, -2.5906, -2.8478, -2.8118, -2.9325, -2.7694]])
        torch.testing.assert_allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)

        EXPECTED_SLICE = {
            7: torch.tensor([-5.8781, -5.8616, -0.1052, -4.7200, -5.8781, -5.8774, -5.8773, -5.8777, -5.8781, -5.8780, -5.8781, -5.8779, -1.0787, 1.7583, -5.8779, -5.8780, -5.8783, -5.8778, -5.8776, -5.8781, -5.8784, -5.8778, -5.8778, -5.8777, -5.8779, -5.8778, -5.8776, -5.8780, -5.8779, -5.8781]),
            8: torch.tensor([-5.8711, -5.8555, -0.1050, -4.7148, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -1.0781, 1.7568, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711]),
        }  # fmt: skip

        print(out[0, 0, :30])
        torch.testing.assert_allclose(
            out[0, 0, :30], EXPECTED_SLICE[self.cuda_compute_capability_major_version], atol=1e-4, rtol=1e-4
        )

        del model
        backend_empty_cache(torch_device)
        gc.collect()
//  [python]
def test_model_7b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = MistralForCausalLM.from_pretrained(
            "mistralai/Mistral-7B-v0.1", device_map="auto", torch_dtype=torch.float16
        )
        input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
        with torch.no_grad():
            out = model(input_ids).logits.cpu()
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-2.5548, -2.5737, -3.0600, -2.5906, -2.8478, -2.8118, -2.9325, -2.7694]])
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)

        EXPECTED_SLICE = {
            7: torch.tensor([-5.8781, -5.8616, -0.1052, -4.7200, -5.8781, -5.8774, -5.8773, -5.8777, -5.8781, -5.8780, -5.8781, -5.8779, -1.0787, 1.7583, -5.8779, -5.8780, -5.8783, -5.8778, -5.8776, -5.8781, -5.8784, -5.8778, -5.8778, -5.8777, -5.8779, -5.8778, -5.8776, -5.8780, -5.8779, -5.8781]),
            8: torch.tensor([-5.8711, -5.8555, -0.1050, -4.7148, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -1.0781, 1.7568, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711, -5.8711]),
        }  # fmt: skip

        print(out[0, 0, :30])
        torch.testing.assert_close(
            out[0, 0, :30], EXPECTED_SLICE[self.cuda_compute_capability_major_version], atol=1e-4, rtol=1e-4
        )

        del model
        backend_empty_cache(torch_device)
        gc.collect()
//  [python]
def test_no_compute(executor, device, dtype):
    def foo(a, b):
        return a, 3.0

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    b = make_tensor((2, 2, 2), device=device, dtype=tdtype)

    thunder_result = traced_foo(a, b=b)
    torch_result = foo(a, b)

    t.assert_close(thunder_result, torch_result)
//  [python]
def test_no_compute(executor, device, dtype):
    def foo(a, b):
        return a, 3.0

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    b = make_tensor((2, 2, 2), device=device, dtype=tdtype)

    thunder_result = traced_foo(a, b=b)
    torch_result = foo(a, b)

    assert_close(thunder_result, torch_result)
//  [python]
def test_partial(executor, device, dtype):
    def foo(a, *, b, c=2):
        return a, b, c

    pfoo = partial(foo, b=3, c=4)
    cpfoo = executor.make_callable(pfoo)

    lc_result = cpfoo(1)
    py_result = pfoo(1)

    t.assert_close(lc_result, py_result)

    # Tests that later partials override earlier partials correctly
    ppfoo = partial(pfoo, b=2, c=8)
    cppfoo = executor.make_callable(ppfoo)

    lc_result = cppfoo(1)
    py_result = ppfoo(1)

    t.assert_close(lc_result, py_result)
//  [python]
def test_partial(executor, device, dtype):
    def foo(a, *, b, c=2):
        return a, b, c

    pfoo = partial(foo, b=3, c=4)
    cpfoo = executor.make_callable(pfoo)

    lc_result = cpfoo(1)
    py_result = pfoo(1)

    assert_close(lc_result, py_result)

    # Tests that later partials override earlier partials correctly
    ppfoo = partial(pfoo, b=2, c=8)
    cppfoo = executor.make_callable(ppfoo)

    lc_result = cppfoo(1)
    py_result = ppfoo(1)

    assert_close(lc_result, py_result)
//  [python]
def test_constant_creation(executor, device, dtype):
    def py_foo(a):
        return a + 1.0

    def foo(a):
        x = prims.convert_element_type(1, float)
        return a + x

    cfoo = thunder.jit(foo, executors=executor.executors_list())

    torch_dtype = ltorch.to_torch_dtype(dtype)
    a = make_tensor((2, 2), device=device, dtype=torch_dtype)

    lc_result = cfoo(a)
    python_result = py_foo(a)

    t.assert_close(lc_result, python_result)
//  [python]
def test_constant_creation(executor, device, dtype):
    def py_foo(a):
        return a + 1.0

    def foo(a):
        x = prims.convert_element_type(1, float)
        return a + x

    cfoo = thunder.jit(foo, executors=executor.executors_list())

    torch_dtype = ltorch.to_torch_dtype(dtype)
    a = make_tensor((2, 2), device=device, dtype=torch_dtype)

    lc_result = cfoo(a)
    python_result = py_foo(a)

    assert_close(lc_result, python_result)
//  [python]
def test_partial_results(executor: TestExecutor, device: str, dtype: dtypes.dtype):
    torch_dtype = ltorch.to_torch_dtype(dtype)
    a = make_tensor((2, 2), device=device, dtype=torch_dtype)

    def foo(a):
        a, b = torch.var_mean(a)
        return a

    cfoo = executor.make_callable(foo, disable_preprocessing=False)

    result = cfoo(a)
    torch_result = foo(a)

    t.assert_close(result, torch_result)
//  [python]
def test_partial_results(executor: TestExecutor, device: str, dtype: dtypes.dtype):
    torch_dtype = ltorch.to_torch_dtype(dtype)
    a = make_tensor((2, 2), device=device, dtype=torch_dtype)

    def foo(a):
        a, b = torch.var_mean(a)
        return a

    cfoo = executor.make_callable(foo, disable_preprocessing=False)

    result = cfoo(a)
    torch_result = foo(a)

    assert_close(result, torch_result)
//  [python]
def test_transforms_vmap_axis_size(executor, device, _):
    from thunder.core.transforms import vmap
    import torch as t

    actual = executor.make_callable_legacy(vmap(lambda: 2, axis_size=4))()
    expected = t.full((4,), 2, device="cpu")
    t.testing.assert_allclose(actual, expected)

    actual = executor.make_callable_legacy(vmap(lambda x: x, axis_size=4))(2)
    t.testing.assert_allclose(actual, expected)
//  [python]
def test_transforms_vmap_axis_size(executor, device, _):
    from thunder.core.transforms import vmap

    actual = executor.make_callable_legacy(vmap(lambda: 2, axis_size=4))()
    expected = torch.full((4,), 2, device="cpu")
    assert_close(actual, expected)

    actual = executor.make_callable_legacy(vmap(lambda x: x, axis_size=4))(2)
    assert_close(actual, expected)
//  [python]
def test_nested_trace(executor, device, _):
    # This test ensures that trace() can be called from within a traced
    # function without leaking the trace context.
    # from thunder import _get_executor

    def foo(a, b):
        return clang.add(a, b)

    def bar(a, b):
        foo_trace = thunder.trace(inline_trace=False)(foo, a, b)
        assert len(foo_trace.bound_symbols) == 4
        assert foo_trace.bound_symbols[-2].sym.name == "add"
        return clang.mul(a, b)

    a = make_tensor((2, 2), device=device, dtype=torch.float32)
    b = make_tensor((2, 2), device=device, dtype=torch.float32)

    bar_trace = thunder.trace()(bar, a, b)
    assert len(bar_trace.bound_symbols) == 4
    assert bar_trace.bound_symbols[-2].sym.name == "mul"

    fusion = executor.make_callable(bar_trace)
    actual = fusion(a, b)
    expected = a * b
    t.assert_close(actual, expected)
//  [python]
def test_nested_trace(executor, device, _):
    # This test ensures that trace() can be called from within a traced
    # function without leaking the trace context.
    # from thunder import _get_executor

    def foo(a, b):
        return clang.add(a, b)

    def bar(a, b):
        foo_trace = thunder.trace(inline_trace=False)(foo, a, b)
        assert len(foo_trace.bound_symbols) == 4
        assert foo_trace.bound_symbols[-2].sym.name == "add"
        return clang.mul(a, b)

    a = make_tensor((2, 2), device=device, dtype=torch.float32)
    b = make_tensor((2, 2), device=device, dtype=torch.float32)

    bar_trace = thunder.trace()(bar, a, b)
    assert len(bar_trace.bound_symbols) == 4
    assert bar_trace.bound_symbols[-2].sym.name == "mul"

    fusion = executor.make_callable(bar_trace)
    actual = fusion(a, b)
    expected = a * b
    assert_close(actual, expected)
//  [python]
def test_eval_trace_duplicate_output(executor, device, _):
    # This test ensures that eval_trace() can evaluate a trace with duplicate
    # outputs.
    from thunder.core.transforms import eval_trace, identity

    def foo1(a):
        return a, a

    a = torch.ones((2, 2), device=device, dtype=torch.float32)

    foo_trace = thunder.trace()(foo1, a)
    assert len(foo_trace.bound_symbols) == 2
    assert foo_trace.bound_symbols[0].sym.name == "unpack_trivial"
    assert len(foo_trace.output) == 2
    assert foo_trace.output[0].name == foo_trace.output[1].name

    def func(a):
        return eval_trace(foo_trace, a)

    actual = executor.make_callable(func)(a)
    t.assert_close(actual, (a, a))

    # Identity is needed to ensure that the duplicated outputs of a symbol
    # don't cause problems.
    def foo2(a):
        a = 1.0 * a
        return a, a

    for foo in [foo1, foo2]:
        foo_trace = thunder.trace()(identity(foo), a)
        assert len(foo_trace.output) == 2
        assert foo_trace.output[0].name == foo_trace.output[1].name
//  [python]
def test_eval_trace_duplicate_output(executor, device, _):
    # This test ensures that eval_trace() can evaluate a trace with duplicate
    # outputs.
    from thunder.core.transforms import eval_trace, identity

    def foo1(a):
        return a, a

    a = torch.ones((2, 2), device=device, dtype=torch.float32)

    foo_trace = thunder.trace()(foo1, a)
    assert len(foo_trace.bound_symbols) == 2
    assert foo_trace.bound_symbols[0].sym.name == "unpack_trivial"
    assert len(foo_trace.output) == 2
    assert foo_trace.output[0].name == foo_trace.output[1].name

    def func(a):
        return eval_trace(foo_trace, a)

    actual = executor.make_callable(func)(a)
    assert_close(actual, (a, a))

    # Identity is needed to ensure that the duplicated outputs of a symbol
    # don't cause problems.
    def foo2(a):
        a = 1.0 * a
        return a, a

    for foo in [foo1, foo2]:
        foo_trace = thunder.trace()(identity(foo), a)
        assert len(foo_trace.output) == 2
        assert foo_trace.output[0].name == foo_trace.output[1].name
//  [python]
def test_normalized_args_prims_sum(executor, device: str, dtype: dtypes.dtype):
    # This test verifies that the recorded trace for a call to prims.sum
    # has its positional and keyword arguments normalized to the same form.
    # See issue "vmap of sum doesn't work when dims are passed as a keyword
    # argument"
    a = make_tensor((2, 2), device=device, dtype=ltorch.to_torch_dtype(dtype))

    def func_dim_posarg(x):
        return prims.sum(x, (0, 1))

    def func_dim_kwarg(x):
        return prims.sum(x, dims=(0, 1))

    c1 = executor.make_callable(func_dim_posarg)
    c2 = executor.make_callable(func_dim_kwarg)
    out1 = c1(a)
    out2 = c2(a)
    t.assert_close(out1, out2)

    trace1 = thunder.last_traces(c1)[0]
    trace2 = thunder.last_traces(c2)[0]
    sum1 = next(s for s in trace1.bound_symbols if s.sym.name == "sum")
    sum2 = next(s for s in trace2.bound_symbols if s.sym.name == "sum")

    assert len(sum1.args) == 2
    assert len(sum2.args) == 2
    assert len(sum1.kwargs) == 0
    assert len(sum2.kwargs) == 0
//  [python]
def test_normalized_args_prims_sum(executor, device: str, dtype: dtypes.dtype):
    # This test verifies that the recorded trace for a call to prims.sum
    # has its positional and keyword arguments normalized to the same form.
    # See issue "vmap of sum doesn't work when dims are passed as a keyword
    # argument"
    a = make_tensor((2, 2), device=device, dtype=ltorch.to_torch_dtype(dtype))

    def func_dim_posarg(x):
        return prims.sum(x, (0, 1))

    def func_dim_kwarg(x):
        return prims.sum(x, dims=(0, 1))

    c1 = executor.make_callable(func_dim_posarg)
    c2 = executor.make_callable(func_dim_kwarg)
    out1 = c1(a)
    out2 = c2(a)
    torch.testing.assert_close(out1, out2)

    trace1 = thunder.last_traces(c1)[0]
    trace2 = thunder.last_traces(c2)[0]
    sum1 = next(s for s in trace1.bound_symbols if s.sym.name == "sum")
    sum2 = next(s for s in trace2.bound_symbols if s.sym.name == "sum")

    assert len(sum1.args) == 2
    assert len(sum2.args) == 2
    assert len(sum1.kwargs) == 0
    assert len(sum2.kwargs) == 0
//  [python]
def test_transforms_vjp_1_2(executor, device, _):
    from thunder.core.transforms import vjp

    # 1 input, 2 outputs
    def func_1_2(x):
        a = clang.sin(x)
        b = clang.add(0.2, a)
        c = clang.asin(b)
        return b, c

    a = make_tensor((2, 3), device=device, dtype=torch.float32)

    g1 = make_tensor((2, 3), device=device, dtype=torch.float32)
    g2 = make_tensor((2, 3), device=device, dtype=torch.float32)

    vjp_eager = executor.make_callable_legacy(vjp(func_1_2))

    primals = (a,)
    cotangents = (g1, g2)
    out_p, grads = vjp_eager(primals, cotangents)
    expected_out_p = executor.make_callable_legacy(func_1_2)(a)
    t.assert_close(out_p, expected_out_p, equal_nan=True, atol=1e-3, rtol=1e-5)

    # Now check the gradients
    # TODO: We will have this automatically tested with OpInfo tests
    aa = a.clone().requires_grad_(True)

    def pt_func_1_2(x):
        a = torch.sin(x)
        b = torch.add(0.2, a)
        c = torch.asin(b)
        return b, c

    out = pt_func_1_2(aa)
    expected_grads = torch.autograd.grad(out, aa, grad_outputs=(g1, g2), retain_graph=True)
    t.assert_close(expected_grads, grads, equal_nan=True, atol=1e-3, rtol=1e-5)
//  [python]
def test_transforms_vjp_1_2(executor, device, _):
    from thunder.core.transforms import vjp

    # 1 input, 2 outputs
    def func_1_2(x):
        a = clang.sin(x)
        b = clang.add(0.2, a)
        c = clang.asin(b)
        return b, c

    a = make_tensor((2, 3), device=device, dtype=torch.float32)

    g1 = make_tensor((2, 3), device=device, dtype=torch.float32)
    g2 = make_tensor((2, 3), device=device, dtype=torch.float32)

    vjp_eager = executor.make_callable_legacy(vjp(func_1_2))

    primals = (a,)
    cotangents = (g1, g2)
    out_p, grads = vjp_eager(primals, cotangents)
    expected_out_p = executor.make_callable_legacy(func_1_2)(a)
    assert_close(out_p, expected_out_p, equal_nan=True, atol=1e-3, rtol=1e-5)

    # Now check the gradients
    # TODO: We will have this automatically tested with OpInfo tests
    aa = a.clone().requires_grad_(True)

    def pt_func_1_2(x):
        a = torch.sin(x)
        b = torch.add(0.2, a)
        c = torch.asin(b)
        return b, c

    out = pt_func_1_2(aa)
    expected_grads = torch.autograd.grad(out, aa, grad_outputs=(g1, g2), retain_graph=True)
    assert_close(expected_grads, grads, equal_nan=True, atol=1e-3, rtol=1e-5)
//  [python]
def test_transforms_inline_jvp_inline_vmap(executor, device, _):
    pytest.xfail("AttributeError: 'NoneType' object has no attribute 'mul'")
    from thunder.core.transforms import vmap, jvp, inline

    if executor == nvFuserExecutor:
        # Couldn't find metadata for 1.0 of type <class 'float'>
        pytest.xfail("Something is failing with the nvFuser executor")

    def func(a, b):
        assert isinstance(a, proxies.TensorProxy)
        assert isinstance(b, proxies.TensorProxy)
        assert a.ndim == 1
        assert a.shape == b.shape
        c = clang.sin(a)
        return clang.mul(clang.add(c, b), 1)

    a = torch.ones(2, 3, device=device, dtype=torch.float32)
    b = torch.ones(2, 3, device=device, dtype=torch.float32) * 2

    args = (a, b)
    out_p, out_t = executor.make_callable(jvp(vmap(func)))(args, args)
    expected_out_p = torch.sin(a) + b
    torch.testing.assert_allclose(out_p, expected_out_p)
    expected_out_t = torch.cos(a) + b
    torch.testing.assert_allclose(out_t, expected_out_t)
//  [python]
def test_transforms_inline_jvp_inline_vmap(executor, device, _):
    pytest.xfail("AttributeError: 'NoneType' object has no attribute 'mul'")
    from thunder.core.transforms import vmap, jvp, inline

    if executor == nvFuserExecutor:
        # Couldn't find metadata for 1.0 of type <class 'float'>
        pytest.xfail("Something is failing with the nvFuser executor")

    def func(a, b):
        assert isinstance(a, proxies.TensorProxy)
        assert isinstance(b, proxies.TensorProxy)
        assert a.ndim == 1
        assert a.shape == b.shape
        c = clang.sin(a)
        return clang.mul(clang.add(c, b), 1)

    a = torch.ones(2, 3, device=device, dtype=torch.float32)
    b = torch.ones(2, 3, device=device, dtype=torch.float32) * 2

    args = (a, b)
    out_p, out_t = executor.make_callable(jvp(vmap(func)))(args, args)
    expected_out_p = torch.sin(a) + b
    assert_close(out_p, expected_out_p)
    expected_out_t = torch.cos(a) + b
    assert_close(out_t, expected_out_t)
//  [python]
def test_torch_tensor_to_memory_format(executor: TestExecutor, device: str, _):
    inp = torch.randn(2, 4, 5, 3, device=device, dtype=torch.float32)

    def torch_to(a, memory_format):
        return a.to(memory_format=memory_format)

    cfn = executor.make_callable(torch_to, disable_preprocessing=False)

    for m_format in [torch.contiguous_format, torch.channels_last, torch.preserve_format]:
        thunder_result = cfn(inp, torch.contiguous_format)
        torch_result = torch_to(inp, torch.contiguous_format)
        torch.testing.assert_allclose(torch_result, thunder_result, check_stride=True)
//  [python]
def test_torch_tensor_to_memory_format(executor: TestExecutor, device: str, _):
    inp = torch.randn(2, 4, 5, 3, device=device, dtype=torch.float32)

    def torch_to(a, memory_format):
        return a.to(memory_format=memory_format)

    cfn = executor.make_callable(torch_to, disable_preprocessing=False)

    for m_format in [torch.contiguous_format, torch.channels_last, torch.preserve_format]:
        thunder_result = cfn(inp, torch.contiguous_format)
        torch_result = torch_to(inp, torch.contiguous_format)
        assert_close(torch_result, thunder_result, check_stride=True)
//  [python]
def test_transforms_inline_vmap_inline_jvp(executor, device, _):
    from thunder.core.transforms import vmap, jvp

    def func(a, b):
        assert isinstance(a, proxies.TensorProxy)
        assert isinstance(b, proxies.TensorProxy)
        assert a.ndim == 1
        assert a.shape == b.shape
        c = clang.sin(a)
        return clang.mul(clang.add(c, b), 1)

    a = torch.ones(2, 3, device=device, dtype=torch.float32)
    b = torch.ones(2, 3, device=device, dtype=torch.float32) * 2

    args = (a, b)
    out_p, out_t = executor.make_callable_legacy(vmap(jvp(func), out_dims=(0, 0)))(args, args)
    expected_out_p = torch.sin(a) + b
    torch.testing.assert_allclose(out_p, expected_out_p)
    expected_out_t = torch.cos(a) + b
    torch.testing.assert_allclose(out_t, expected_out_t)
//  [python]
def test_transforms_inline_vmap_inline_jvp(executor, device, _):
    from thunder.core.transforms import vmap, jvp

    def func(a, b):
        assert isinstance(a, proxies.TensorProxy)
        assert isinstance(b, proxies.TensorProxy)
        assert a.ndim == 1
        assert a.shape == b.shape
        c = clang.sin(a)
        return clang.mul(clang.add(c, b), 1)

    a = torch.ones(2, 3, device=device, dtype=torch.float32)
    b = torch.ones(2, 3, device=device, dtype=torch.float32) * 2

    args = (a, b)
    out_p, out_t = executor.make_callable_legacy(vmap(jvp(func), out_dims=(0, 0)))(args, args)
    expected_out_p = torch.sin(a) + b
    assert_close(out_p, expected_out_p)
    expected_out_t = torch.cos(a) + b
    assert_close(out_t, expected_out_t)
//  [python]
def test_transforms_vjp_2_1(executor, device, _):
    from thunder.core.transforms import vjp

    def pt_func_2_1(x, y):
        a = torch.sin(x + y)
        b = torch.add(0.2, a)
        c = torch.asin(b)
        return c

    def func_2_1(x, y):
        a = clang.sin(x + y)
        b = clang.add(0.2, a)
        c = clang.asin(b)
        return c

    vjp_eager = executor.make_callable_legacy(vjp(func_2_1))
    a = make_tensor((2, 3), device=device, dtype=torch.float32)
    b = make_tensor((2, 3), device=device, dtype=torch.float32)
    g1 = make_tensor((2, 3), device=device, dtype=torch.float32)
    primals = (a, b)
    cotangents = (g1,)
    out_p, grads = vjp_eager(primals, cotangents)
    expected_out_p = executor.make_callable_legacy(func_2_1)(*primals)
    t.assert_close(out_p, expected_out_p, equal_nan=True)

    aa = a.clone().requires_grad_(True)
    bb = b.clone().requires_grad_(True)
    out = pt_func_2_1(aa, bb)
    expected_grads = torch.autograd.grad(out, [aa, bb], grad_outputs=(g1,), retain_graph=True)
    t.assert_close(expected_grads, grads, equal_nan=True)
//  [python]
def test_transforms_vjp_2_1(executor, device, _):
    from thunder.core.transforms import vjp

    def pt_func_2_1(x, y):
        a = torch.sin(x + y)
        b = torch.add(0.2, a)
        c = torch.asin(b)
        return c

    def func_2_1(x, y):
        a = clang.sin(x + y)
        b = clang.add(0.2, a)
        c = clang.asin(b)
        return c

    vjp_eager = executor.make_callable_legacy(vjp(func_2_1))
    a = make_tensor((2, 3), device=device, dtype=torch.float32)
    b = make_tensor((2, 3), device=device, dtype=torch.float32)
    g1 = make_tensor((2, 3), device=device, dtype=torch.float32)
    primals = (a, b)
    cotangents = (g1,)
    out_p, grads = vjp_eager(primals, cotangents)
    expected_out_p = executor.make_callable_legacy(func_2_1)(*primals)
    assert_close(out_p, expected_out_p, equal_nan=True)

    aa = a.clone().requires_grad_(True)
    bb = b.clone().requires_grad_(True)
    out = pt_func_2_1(aa, bb)
    expected_grads = torch.autograd.grad(out, [aa, bb], grad_outputs=(g1,), retain_graph=True)
    assert_close(expected_grads, grads, equal_nan=True)
//  [python]
def test_transforms_vjp_2_2_kwarg(executor, device, _):
    # This test ensures that combination of positional and keyword arguments
    # is differentiable.
    from thunder.core.transforms import vjp

    # 2 inputs, 1 kwarg, 2 outputs
    def func_2_2(x, y, *, z):
        def func(x):
            a = clang.sin(x)
            b = clang.add(0.2, a)
            c = clang.asin(b)
            return c

        a, b = func(x), func(y)
        c = clang.add(a, b)
        d = clang.add(c, func(z))
        return c, d

    x = make_tensor((2, 3), device=device, dtype=torch.float64)
    y = make_tensor((2, 3), device=device, dtype=torch.float64)
    z = make_tensor((2, 3), device=device, dtype=torch.float64)

    g1 = make_tensor((2, 3), device=device, dtype=torch.float64)
    g2 = make_tensor((2, 3), device=device, dtype=torch.float64)

    vjp_eager = executor.make_callable_legacy(vjp(func_2_2))

    primals = (x, y)
    primal_kwargs = {"z": z}
    cotangents = (g1, g2)
    out_p, grads = vjp_eager(primals, cotangents, **primal_kwargs)
    expected_out_p = executor.make_callable(func_2_2)(*primals, **primal_kwargs)
    t.assert_close(out_p, expected_out_p, equal_nan=True)

    # Now check the gradients
    # TODO: We will have this automatically tested with OpInfo tests
    xx = x.clone().requires_grad_(True)
    yy = y.clone().requires_grad_(True)
    zz = z.clone().requires_grad_(True)

    def pt_func_2_2(x, y, *, z):
        def func(x):
            a = torch.sin(x)
            b = torch.add(0.2, a)
            c = torch.asin(b)
            return c

        a, b = func(x), func(y)
        c = torch.add(a, b)
        d = torch.add(c, func(z))
        return c, d

    out = pt_func_2_2(xx, yy, z=zz)
    expected_grads = torch.autograd.grad(out, [xx, yy, zz], grad_outputs=(g1, g2), retain_graph=True)
    # vjp returns a tuple of (primals, cotangents) where cotangents is a tuple of
    # derivatives with respect to the positional arguments and a dict of derivatives
    # with respect to the keyword arguments.
    *gprimals, gkwargs = grads
    t.assert_close(expected_grads[:2], gprimals, equal_nan=True)
    t.assert_close(expected_grads[2], gkwargs["z"], equal_nan=True)
//  [python]
def test_transforms_vjp_2_2_kwarg(executor, device, _):
    # This test ensures that combination of positional and keyword arguments
    # is differentiable.
    from thunder.core.transforms import vjp

    # 2 inputs, 1 kwarg, 2 outputs
    def func_2_2(x, y, *, z):
        def func(x):
            a = clang.sin(x)
            b = clang.add(0.2, a)
            c = clang.asin(b)
            return c

        a, b = func(x), func(y)
        c = clang.add(a, b)
        d = clang.add(c, func(z))
        return c, d

    x = make_tensor((2, 3), device=device, dtype=torch.float64)
    y = make_tensor((2, 3), device=device, dtype=torch.float64)
    z = make_tensor((2, 3), device=device, dtype=torch.float64)

    g1 = make_tensor((2, 3), device=device, dtype=torch.float64)
    g2 = make_tensor((2, 3), device=device, dtype=torch.float64)

    vjp_eager = executor.make_callable_legacy(vjp(func_2_2))

    primals = (x, y)
    primal_kwargs = {"z": z}
    cotangents = (g1, g2)
    out_p, grads = vjp_eager(primals, cotangents, **primal_kwargs)
    expected_out_p = executor.make_callable(func_2_2)(*primals, **primal_kwargs)
    assert_close(out_p, expected_out_p, equal_nan=True)

    # Now check the gradients
    # TODO: We will have this automatically tested with OpInfo tests
    xx = x.clone().requires_grad_(True)
    yy = y.clone().requires_grad_(True)
    zz = z.clone().requires_grad_(True)

    def pt_func_2_2(x, y, *, z):
        def func(x):
            a = torch.sin(x)
            b = torch.add(0.2, a)
            c = torch.asin(b)
            return c

        a, b = func(x), func(y)
        c = torch.add(a, b)
        d = torch.add(c, func(z))
        return c, d

    out = pt_func_2_2(xx, yy, z=zz)
    expected_grads = torch.autograd.grad(out, [xx, yy, zz], grad_outputs=(g1, g2), retain_graph=True)
    # vjp returns a tuple of (primals, cotangents) where cotangents is a tuple of
    # derivatives with respect to the positional arguments and a dict of derivatives
    # with respect to the keyword arguments.
    *gprimals, gkwargs = grads
    assert_close(expected_grads[:2], gprimals, equal_nan=True)
    assert_close(expected_grads[2], gkwargs["z"], equal_nan=True)
//  [python]
def test_torch_scaled_dot_product_attention_non_decomposed(executor, device, _):
    n_embd = 32
    B = 2
    qkv = make_tensor(B, n_embd, 3 * n_embd, device=device, dtype=torch.float32)

    def func(qkv):
        # Preprocessing doesn't support nonlocal variables yet, so
        # we need to define the constants here.
        n_embd = 32
        n_head = 16
        B = 2
        T = 32
        C = n_embd
        q, k, v = qkv.split(n_embd, dim=2)  # Results in 3 non-contiguous but "viewable" tensors
        k = k.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
        y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0, is_causal=True)
        return y

    compiled = thunder.jit(func, executors=executor.executors_list())
    out = compiled(qkv)
    traces = thunder.last_traces(compiled)
    torch.testing.assert_allclose(out, func(qkv))
    assert "scaled_dot_product_attention" in tuple(bsym.sym.id for bsym in traces[-1].bound_symbols)
//  [python]
def test_torch_scaled_dot_product_attention_non_decomposed(executor, device, _):
    n_embd = 32
    B = 2
    qkv = make_tensor(B, n_embd, 3 * n_embd, device=device, dtype=torch.float32)

    def func(qkv):
        # Preprocessing doesn't support nonlocal variables yet, so
        # we need to define the constants here.
        n_embd = 32
        n_head = 16
        B = 2
        T = 32
        C = n_embd
        q, k, v = qkv.split(n_embd, dim=2)  # Results in 3 non-contiguous but "viewable" tensors
        k = k.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
        y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0, is_causal=True)
        return y

    compiled = thunder.jit(func, executors=executor.executors_list())
    out = compiled(qkv)
    traces = thunder.last_traces(compiled)
    torch.testing.assert_close(out, func(qkv))
    assert "scaled_dot_product_attention" in tuple(bsym.sym.id for bsym in traces[-1].bound_symbols)
//  [python]
schedule = ScheduleGPipe(stage, chunks, loss_fn=loss_fn)

# Input data
x = torch.randn(batch_size, in_dim, device=device)
target = torch.randn(batch_size, out_dim, device=device)

# Run the pipeline with input `x`. Divide the batch into 4 micro-batches
# and run them in parallel on the pipeline
if rank == 0:
    schedule.step(x)
elif rank == world_size - 1:
    losses = []
    output = schedule.step(target=target, losses=losses)
else:
    schedule.step()

if rank == world_size - 1:
    # Run the original code and get the output for comparison
    reference_output = mn(x)
    # Compare numerics of pipeline and original model
    torch.testing.assert_allclose(output, reference_output)
    print(f"Loss of microbatches: {losses}")
    print(" Pipeline parallel model ran successfully! ".center(80, "*"))
//  [python]
schedule = ScheduleGPipe(stage, chunks, loss_fn=loss_fn)

# Input data
x = torch.randn(batch_size, in_dim, device=device)
target = torch.randn(batch_size, out_dim, device=device)

# Run the pipeline with input `x`. Divide the batch into 4 micro-batches
# and run them in parallel on the pipeline
if rank == 0:
    schedule.step(x)
elif rank == world_size - 1:
    losses = []
    output = schedule.step(target=target, losses=losses)
else:
    schedule.step()

if rank == world_size - 1:
    # Run the original code and get the output for comparison
    reference_output = mn(x)
    # Compare numerics of pipeline and original model
    torch.testing.assert_close(output, reference_output)
    print(f"Loss of microbatches: {losses}")
    print(" Pipeline parallel model ran successfully! ".center(80, "*"))
//  [python]
def test_linear_dynamic(dtype, bias, in_features, out_features, N):
    with torch.no_grad():
        m = LinearModule(in_features, out_features, bias=bias).eval()
        m_q = torch.quantization.quantize_dynamic(m, {torch.nn.Linear},
                                                  dtype=torch.qint8)
        x = torch.randn(N, in_features)
        out = m_q(x)
        out = out.cuda().to(dtype=dtype)

        m_cuda = m.cuda().to(dtype=dtype)
        m_q_cuda = torch.quantization.quantize_dynamic(
            m_cuda, {torch.nn.Linear},
            dtype=torch.qint8).to(dtype=dtype)
        x_cuda = x.cuda().to(dtype=dtype)
        out_cuda = m_q_cuda(x_cuda)

        torch.testing.assert_allclose(out_cuda, out, rtol=3e-2, atol=3e-2)
//  [python]
def test_linear_dynamic(dtype, bias, in_features, out_features, N):
    with torch.no_grad():
        m = LinearModule(in_features, out_features, bias=bias).eval()
        m_q = torch.quantization.quantize_dynamic(m, {torch.nn.Linear},
                                                  dtype=torch.qint8)
        x = torch.randn(N, in_features)
        out = m_q(x)
        out = out.cuda().to(dtype=dtype)

        m_cuda = m.cuda().to(dtype=dtype)
        m_q_cuda = torch.quantization.quantize_dynamic(
            m_cuda, {torch.nn.Linear},
            dtype=torch.qint8).to(dtype=dtype)
        x_cuda = x.cuda().to(dtype=dtype)
        out_cuda = m_q_cuda(x_cuda)

        torch.testing.assert_close(out_cuda, out, rtol=3e-2, atol=3e-2)
//  [python]
def test_load_full_checkpoint_only_model(tmp_path):
    strategy = ThunderFSDPStrategy()
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy)
    fabric.launch()

    checkpoint_path = tmp_path / "foo"
    checkpoint_path = fabric.broadcast(checkpoint_path)
    if fabric.global_rank == 0:
        model = MyModel(4)
        expected = model.state_dict()
        torch.save(expected, checkpoint_path)
    fabric.barrier()
    expected = torch.load(checkpoint_path)

    # before sharding
    model = MyModel(4)
    fabric.load_raw(checkpoint_path, model)
    torch.testing.assert_allclose(model.state_dict(), expected)

    # after sharding
    model = MyModel(4)
    model = fabric.setup(model)
    fabric.load_raw(checkpoint_path, model)
    from thunder.distributed import _unshard_params

    # unshard this model's parameters to compare with the original state dict before sharding
    _unshard_params(model, model.process_group_for_ddp, True)
    actual = model.state_dict()
    # `_unshard_params` doesnt offload buffers at the moment
    assert actual["buf"].device.type == "cuda"
    actual["buf"] = actual["buf"].to(device="cpu")
    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_load_full_checkpoint_only_model(tmp_path):
    strategy = ThunderFSDPStrategy()
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy)
    fabric.launch()

    checkpoint_path = tmp_path / "foo"
    checkpoint_path = fabric.broadcast(checkpoint_path)
    if fabric.global_rank == 0:
        model = MyModel(4)
        expected = model.state_dict()
        torch.save(expected, checkpoint_path)
    fabric.barrier()
    expected = torch.load(checkpoint_path)

    # before sharding
    model = MyModel(4)
    fabric.load_raw(checkpoint_path, model)
    torch.testing.assert_close(model.state_dict(), expected)

    # after sharding
    model = MyModel(4)
    model = fabric.setup(model)
    fabric.load_raw(checkpoint_path, model)
    from thunder.distributed import _unshard_params

    # unshard this model's parameters to compare with the original state dict before sharding
    _unshard_params(model, model.process_group_for_ddp, True)
    actual = model.state_dict()
    # `_unshard_params` doesnt offload buffers at the moment
    assert actual["buf"].device.type == "cuda"
    actual["buf"] = actual["buf"].to(device="cpu")
    torch.testing.assert_close(actual, expected)
//  [python]
def test_save_load_full_checkpoint(tmp_path):
    strategy = ThunderFSDPStrategy(state_dict_type="full", broadcast_from=0)
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy)
    fabric.launch()

    model = MyModel(4)
    expected = model.state_dict()

    # save a sharded model
    model = fabric.setup(model)
    state = {"model": model, "stateful": StatefulThing(), "primitive": 123}
    checkpoint_path = tmp_path / "foo"
    fabric.save(checkpoint_path, state)

    # assert the file contents
    if fabric.global_rank == 0:
        checkpoint = torch.load(checkpoint_path)
        # cpu_offload is enabled by default
        assert checkpoint == {
            "model": {
                "buf": TensorLike("cpu", tuple()),
                "inner.l.weight": TensorLike("cpu", (8, 4)),
                "l.bias": TensorLike("cpu", (4,)),
                "l.weight": TensorLike("cpu", (4, 2)),
            },
            "stateful": {"thing": 1},
            "primitive": 123,
        }
        t.assert_close(checkpoint["model"], expected)

    # load its weights into a different sharded model
    model = MyModel(4)
    model = fabric.setup(model)
    state = {"model": model, "stateful": StatefulThing(), "primitive": 321}
    fabric.load(checkpoint_path, state)

    from thunder.distributed import _unshard_params

    # unshard this model's parameters to compare with the original state dict before sharding
    _unshard_params(model, model.process_group_for_ddp, True)
    # we loaded rank 0's weights, so this would fail in the other ranks
    if fabric.global_rank == 0:
        actual = model.state_dict()
        # `_unshard_params` doesnt offload buffers at the moment
        assert actual["buf"].device.type == "cuda"
        actual["buf"] = actual["buf"].to(device="cpu")
        t.assert_close(actual, expected)
    assert state["primitive"] == 123
//  [python]
def test_save_load_full_checkpoint(tmp_path):
    strategy = ThunderFSDPStrategy(state_dict_type="full", broadcast_from=0)
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy)
    fabric.launch()

    model = MyModel(4)
    expected = model.state_dict()

    # save a sharded model
    model = fabric.setup(model)
    state = {"model": model, "stateful": StatefulThing(), "primitive": 123}
    checkpoint_path = tmp_path / "foo"
    fabric.save(checkpoint_path, state)

    # assert the file contents
    if fabric.global_rank == 0:
        checkpoint = torch.load(checkpoint_path)
        # cpu_offload is enabled by default
        assert checkpoint == {
            "model": {
                "buf": TensorLike("cpu", tuple()),
                "inner.l.weight": TensorLike("cpu", (8, 4)),
                "l.bias": TensorLike("cpu", (4,)),
                "l.weight": TensorLike("cpu", (4, 2)),
            },
            "stateful": {"thing": 1},
            "primitive": 123,
        }
        torch.testing.assert_close(checkpoint["model"], expected)

    # load its weights into a different sharded model
    model = MyModel(4)
    model = fabric.setup(model)
    state = {"model": model, "stateful": StatefulThing(), "primitive": 321}
    fabric.load(checkpoint_path, state)

    from thunder.distributed import _unshard_params

    # unshard this model's parameters to compare with the original state dict before sharding
    _unshard_params(model, model.process_group_for_ddp, True)
    # we loaded rank 0's weights, so this would fail in the other ranks
    if fabric.global_rank == 0:
        actual = model.state_dict()
        # `_unshard_params` doesnt offload buffers at the moment
        assert actual["buf"].device.type == "cuda"
        actual["buf"] = actual["buf"].to(device="cpu")
        torch.testing.assert_close(actual, expected)
    assert state["primitive"] == 123
//  [python]
def test_save_load_sharded_checkpoint(tmp_path):
    pytest.skip("Temporarily disabled, often exceeds 5 min timeout")

    strategy = ThunderFSDPStrategy(state_dict_type="sharded", broadcast_from=0)
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy)
    fabric.launch()

    model = MyModel(4)
    expected = model.state_dict()

    # save a sharded model
    model = fabric.setup(model)
    state = {"model": model, "stateful": StatefulThing(), "primitive": 123}
    fabric.save(tmp_path, state)

    # assert the file contents
    if fabric.global_rank == 0:
        assert set(os.listdir(tmp_path)) == {"meta.pt", "__1_0.distcp", "__0_0.distcp", ".metadata"}

        metadata = torch.load(tmp_path / "meta.pt")
        assert metadata == {"stateful": {"thing": 1}, "primitive": 123}

        checkpoint = distributed_ckpt_to_regular(tmp_path)
        # cpu_offload is enabled by default
        assert checkpoint == {
            "model": {
                "buf": TensorLike("cpu", tuple()),
                "inner.l.weight": TensorLike("cpu", (8, 4)),
                "l.bias": TensorLike("cpu", (4,)),
                "l.weight": TensorLike("cpu", (4, 2)),
            }
        }
        t.assert_close(checkpoint["model"], expected)

    # load its weights into a different sharded model
    model = MyModel(4)
    model = fabric.setup(model)
    state = {"model": model, "stateful": StatefulThing(), "primitive": 321}
    fabric.load(tmp_path, state)

    from thunder.distributed import _unshard_params

    # unshard this model's parameters to compare with the original state dict before sharding
    _unshard_params(model, model.process_group_for_ddp, True)
    # we loaded rank 0's weights, so this would fail in the other ranks
    if fabric.global_rank == 0:
        actual = model.state_dict()
        # `_unshard_params` doesnt offload buffers at the moment
        assert actual["buf"].device.type == "cuda"
        actual["buf"] = actual["buf"].to(device="cpu")
        t.assert_close(actual, expected)
    assert state["primitive"] == 123
//  [python]
def test_save_load_sharded_checkpoint(tmp_path):
    pytest.skip("Temporarily disabled, often exceeds 5 min timeout")

    strategy = ThunderFSDPStrategy(state_dict_type="sharded", broadcast_from=0)
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy)
    fabric.launch()

    model = MyModel(4)
    expected = model.state_dict()

    # save a sharded model
    model = fabric.setup(model)
    state = {"model": model, "stateful": StatefulThing(), "primitive": 123}
    fabric.save(tmp_path, state)

    # assert the file contents
    if fabric.global_rank == 0:
        assert set(os.listdir(tmp_path)) == {"meta.pt", "__1_0.distcp", "__0_0.distcp", ".metadata"}

        metadata = torch.load(tmp_path / "meta.pt")
        assert metadata == {"stateful": {"thing": 1}, "primitive": 123}

        checkpoint = distributed_ckpt_to_regular(tmp_path)
        # cpu_offload is enabled by default
        assert checkpoint == {
            "model": {
                "buf": TensorLike("cpu", tuple()),
                "inner.l.weight": TensorLike("cpu", (8, 4)),
                "l.bias": TensorLike("cpu", (4,)),
                "l.weight": TensorLike("cpu", (4, 2)),
            }
        }
        torch.testing.assert_close(checkpoint["model"], expected)

    # load its weights into a different sharded model
    model = MyModel(4)
    model = fabric.setup(model)
    state = {"model": model, "stateful": StatefulThing(), "primitive": 321}
    fabric.load(tmp_path, state)

    from thunder.distributed import _unshard_params

    # unshard this model's parameters to compare with the original state dict before sharding
    _unshard_params(model, model.process_group_for_ddp, True)
    # we loaded rank 0's weights, so this would fail in the other ranks
    if fabric.global_rank == 0:
        actual = model.state_dict()
        # `_unshard_params` doesnt offload buffers at the moment
        assert actual["buf"].device.type == "cuda"
        actual["buf"] = actual["buf"].to(device="cpu")
        torch.testing.assert_close(actual, expected)
    assert state["primitive"] == 123