jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_SpectralConv1D(factorization, implementation):
    """Test for SpectralConv1D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    """
    conv = SpectralConv(
        10, 11, (5,), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )

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

    x = torch.randn(2, 10, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    torch.testing.assert_close(res_dense, res)
//  [python]
def test_cosine_sim_compare():
    a = torch.rand(128)
    b = torch.rand(128)
    torch.testing.assert_allclose(
        torch_metrics.cosine_sim(a, b),
        torch.from_numpy(np_metrics.cosine_sim(a.numpy(), b.numpy())),
    )

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    torch.testing.assert_allclose(
        torch_metrics.cosine_sim(a, b),
        torch.from_numpy(np_metrics.cosine_sim(a.numpy(), b.numpy())),
    )
//  [python]
def test_cosine_sim_compare():
    a = torch.rand(128)
    b = torch.rand(128)
    torch.testing.assert_close(
        torch_metrics.cosine_sim(a, b),
        torch.from_numpy(np_metrics.cosine_sim(a.numpy(), b.numpy())),
    )

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    torch.testing.assert_close(
        torch_metrics.cosine_sim(a, b),
        torch.from_numpy(np_metrics.cosine_sim(a.numpy(), b.numpy())),
    )
//  [python]
def test_euclidean_dist_compare():
    a = torch.rand(128)
    b = torch.rand(128)
    torch.testing.assert_allclose(
        torch_metrics.euclidean_dist(a, b),
        torch.from_numpy(np_metrics.euclidean_dist(a.numpy(), b.numpy())).to(
            torch.float32
        ),
    )

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    torch.testing.assert_allclose(
        torch_metrics.euclidean_dist(a, b),
        torch.from_numpy(np_metrics.euclidean_dist(a.numpy(), b.numpy())),
    )
//  [python]
def test_euclidean_dist_compare():
    a = torch.rand(128)
    b = torch.rand(128)
    torch.testing.assert_close(
        torch_metrics.euclidean_dist(a, b),
        torch.from_numpy(np_metrics.euclidean_dist(a.numpy(), b.numpy())).to(
            torch.float32
        ),
    )

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    torch.testing.assert_close(
        torch_metrics.euclidean_dist(a, b),
        torch.from_numpy(np_metrics.euclidean_dist(a.numpy(), b.numpy())),
    )
//  [python]
def test_xfail_func_with_non_tensor_args(self):
        class Foo(torch.nn.Module):
            def forward(self, x, b=1.0):
                y = x + b
                z = y.relu()
                return (y, z)

        func = Foo()

        tensor_x = torch.randn(1, 1, 2, dtype=torch.float32)

        onnx_program = torch.onnx.dynamo_export(
            func,
            tensor_x,
            8.0,
            export_options=torch.onnx.ExportOptions(
                op_level_debug=self.op_level_debug,
                dynamic_shapes=self.dynamic_shapes,
            ),
        )
        onnx_test_common.assert_dynamic_shapes(onnx_program, self.dynamic_shapes)
        onnx_format_args = onnx_program.adapt_torch_inputs_to_onnx(tensor_x, b=8.0)
        ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(func(tensor_x, 8.0))
        ort_outputs = onnx_test_common.run_ort(onnx_program, onnx_format_args)
        for ref_output, ort_output in zip(ref_outputs, ort_outputs):
            torch.testing.assert_allclose(ref_output, torch.tensor(ort_output))

        # test on different non-tensor input - xfail
        onnx_format_args = onnx_program.adapt_torch_inputs_to_onnx(tensor_x, b=9.0)
        ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(func(tensor_x, 9.0))
        _ = onnx_test_common.run_ort(onnx_program, onnx_format_args)
        for ref_output, ort_output in zip(ref_outputs, ort_outputs):
            torch.testing.assert_allclose(ref_output, torch.tensor(ort_output))
//  [python]
def test_xfail_func_with_non_tensor_args(self):
        class Foo(torch.nn.Module):
            def forward(self, x, b=1.0):
                y = x + b
                z = y.relu()
                return (y, z)

        func = Foo()

        tensor_x = torch.randn(1, 1, 2, dtype=torch.float32)

        onnx_program = torch.onnx.dynamo_export(
            func,
            tensor_x,
            8.0,
            export_options=torch.onnx.ExportOptions(
                op_level_debug=self.op_level_debug,
                dynamic_shapes=self.dynamic_shapes,
            ),
        )
        onnx_test_common.assert_dynamic_shapes(onnx_program, self.dynamic_shapes)
        onnx_format_args = onnx_program.adapt_torch_inputs_to_onnx(tensor_x, b=8.0)
        ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(func(tensor_x, 8.0))
        ort_outputs = onnx_test_common.run_ort(onnx_program, onnx_format_args)
        for ref_output, ort_output in zip(ref_outputs, ort_outputs):
            torch.testing.assert_close(ref_output, torch.tensor(ort_output))

        # test on different non-tensor input - xfail
        onnx_format_args = onnx_program.adapt_torch_inputs_to_onnx(tensor_x, b=9.0)
        ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(func(tensor_x, 9.0))
        _ = onnx_test_common.run_ort(onnx_program, onnx_format_args)
        for ref_output, ort_output in zip(ref_outputs, ort_outputs):
            torch.testing.assert_close(ref_output, torch.tensor(ort_output))
//  [python]
def test_sqeuclidean_dist_compare():
    a = torch.rand(128)
    b = torch.rand(128)
    t.assert_allclose(
        torch_metrics.sqeuclidean_dist(a, b),
        torch.from_numpy(np_metrics.sqeuclidean_dist(a.numpy(), b.numpy())).to(
            torch.float32
        ),
    )

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    t.assert_allclose(
        torch_metrics.sqeuclidean_dist(a, b),
        torch.from_numpy(np_metrics.sqeuclidean_dist(a.numpy(), b.numpy())),
    )
//  [python]
def test_sqeuclidean_dist_compare():
    a = torch.rand(128)
    b = torch.rand(128)
    torch.testing.assert_close(
        torch_metrics.sqeuclidean_dist(a, b),
        torch.from_numpy(np_metrics.sqeuclidean_dist(a.numpy(), b.numpy())).to(
            torch.float32
        ),
    )

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    torch.testing.assert_close(
        torch_metrics.sqeuclidean_dist(a, b),
        torch.from_numpy(np_metrics.sqeuclidean_dist(a.numpy(), b.numpy())),
    )
//  [python]
def test_simple(self) -> None:
        # See N4816561 to view these results graphically
        def f1(x: float) -> float:
            return x

        def f2(x: float) -> float:
            return x * x - 10 * x + 10  # min at x = 5

        def f3(x: float) -> float:
            # bumpy function, overall min at x=30
            return (x - 30) ** 2 + 100 * math.sin(x)

        def f4(x: float) -> float:
            # spiky/non-smooth function, min at x = 30
            return (x - 30) ** 2 + (x % 10) * 100

        results = TestLuusJaakolaSearch.evaluate(0, 100, f1)
        want = torch.tensor([[0, 0], [0, 0], [0, 0], [0, 0], [0, 0]], dtype=torch.int64)
        torch.testing.assert_allclose(results, want)

        results = TestLuusJaakolaSearch.evaluate(0, 100, f2)
        want = torch.tensor(
            [
                [3.51914, -12.80705],
                [4.22958, -14.40646],
                [5.41303, -14.82940],
                [2.35012, -7.97811],
                [4.18552, -14.33662],
            ]
        )
        torch.testing.assert_allclose(results, want)

        results = TestLuusJaakolaSearch.evaluate(0, 100, f3)
        want = torch.tensor(
            [
                [36.58517, -46.37988],
                [29.73184, -99.28705],
                [37.67208, 56.15779],
                [35.85468, -62.00219],
                [41.76223, 58.69744],
            ]
        )
        torch.testing.assert_allclose(results, want)

        results = TestLuusJaakolaSearch.evaluate(0, 100, f4)
        want = torch.tensor(
            [
                [23.68681, 408.53735],
                [31.62534, 165.17535],
                [32.81968, 289.91898],
                [42.81567, 445.80777],
                [22.53002, 308.80225],
            ]
        )
        torch.testing.assert_allclose(results, want)
//  [python]
def test_simple(self) -> None:
        # See N4816561 to view these results graphically
        def f1(x: float) -> float:
            return x

        def f2(x: float) -> float:
            return x * x - 10 * x + 10  # min at x = 5

        def f3(x: float) -> float:
            # bumpy function, overall min at x=30
            return (x - 30) ** 2 + 100 * math.sin(x)

        def f4(x: float) -> float:
            # spiky/non-smooth function, min at x = 30
            return (x - 30) ** 2 + (x % 10) * 100

        results = TestLuusJaakolaSearch.evaluate(0, 100, f1)
        want = torch.tensor([[0, 0], [0, 0], [0, 0], [0, 0], [0, 0]], dtype=torch.int64)
        torch.testing.assert_close(results, want)

        results = TestLuusJaakolaSearch.evaluate(0, 100, f2)
        want = torch.tensor(
            [
                [3.51914, -12.80705],
                [4.22958, -14.40646],
                [5.41303, -14.82940],
                [2.35012, -7.97811],
                [4.18552, -14.33662],
            ]
        )
        torch.testing.assert_close(results, want)

        results = TestLuusJaakolaSearch.evaluate(0, 100, f3)
        want = torch.tensor(
            [
                [36.58517, -46.37988],
                [29.73184, -99.28705],
                [37.67208, 56.15779],
                [35.85468, -62.00219],
                [41.76223, 58.69744],
            ]
        )
        torch.testing.assert_close(results, want)

        results = TestLuusJaakolaSearch.evaluate(0, 100, f4)
        want = torch.tensor(
            [
                [23.68681, 408.53735],
                [31.62534, 165.17535],
                [32.81968, 289.91898],
                [42.81567, 445.80777],
                [22.53002, 308.80225],
            ]
        )
        torch.testing.assert_close(results, want)
//  [python]
def test_SpectralConv(factorization, implementation):
    """Test for SpectralConv of any order
    
    Compares Factorized and Dense convolution output
    Verifies that a dense conv and factorized conv with the same weight produce the same output

    Checks the output size

    Verifies that dynamically changing the number of Fourier modes doesn't break the conv
    """
    modes = (10, 8, 6, 6)
    incremental_modes = (6, 6, 4, 4)

    # Test for Conv1D to Conv4D
    for dim in [1, 2, 3, 4]:
        conv = SpectralConv(
            3, 3, modes[:dim], n_layers=1, bias=False, implementation=implementation, factorization=factorization)

        conv_dense = SpectralConv(
            3, 3, modes[:dim], n_layers=1, bias=False, implementation='reconstructed', factorization=None)

        for i in range(2**(dim-1)):
            conv_dense.weight[i] = FactorizedTensor.from_tensor(conv.weight[i].to_tensor(), rank=None, factorization='ComplexDense')

        x = torch.randn(2, 3, *(12, )*dim)

        res_dense = conv_dense(x)
        res = conv(x)
        res_shape = res.shape

        torch.testing.assert_allclose(res_dense, res)

        # Dynamically reduce the number of modes in Fourier space
        conv.incremental_n_modes = incremental_modes[:dim]
        res = conv(x)
        assert res_shape == res.shape

        # Downsample outputs
        block = SpectralConv(
            3, 4, modes[:dim], n_layers=1, output_scaling_factor=0.5)
    
        x = torch.randn(2, 3, *(12, )*dim)
        res = block(x)
        assert(list(res.shape[2:]) == [12//2]*dim)
        
        # Upsample outputs
        block = SpectralConv(
            3, 4, modes[:dim], n_layers=1, output_scaling_factor=2)
    
        x = torch.randn(2, 3, *(12, )*dim)
        res = block(x)
        assert res.shape[1] == 4 # Check out channels
        assert(list(res.shape[2:]) == [12*2]*dim)
//  [python]
def test_SpectralConv(factorization, implementation):
    """Test for SpectralConv of any order
    
    Compares Factorized and Dense convolution output
    Verifies that a dense conv and factorized conv with the same weight produce the same output

    Checks the output size

    Verifies that dynamically changing the number of Fourier modes doesn't break the conv
    """
    modes = (10, 8, 6, 6)
    incremental_modes = (6, 6, 4, 4)

    # Test for Conv1D to Conv4D
    for dim in [1, 2, 3, 4]:
        conv = SpectralConv(
            3, 3, modes[:dim], n_layers=1, bias=False, implementation=implementation, factorization=factorization)

        conv_dense = SpectralConv(
            3, 3, modes[:dim], n_layers=1, bias=False, implementation='reconstructed', factorization=None)

        for i in range(2**(dim-1)):
            conv_dense.weight[i] = FactorizedTensor.from_tensor(conv.weight[i].to_tensor(), rank=None, factorization='ComplexDense')

        x = torch.randn(2, 3, *(12, )*dim)

        res_dense = conv_dense(x)
        res = conv(x)
        res_shape = res.shape

        torch.testing.assert_close(res_dense, res)

        # Dynamically reduce the number of modes in Fourier space
        conv.incremental_n_modes = incremental_modes[:dim]
        res = conv(x)
        assert res_shape == res.shape

        # Downsample outputs
        block = SpectralConv(
            3, 4, modes[:dim], n_layers=1, output_scaling_factor=0.5)
    
        x = torch.randn(2, 3, *(12, )*dim)
        res = block(x)
        assert(list(res.shape[2:]) == [12//2]*dim)
        
        # Upsample outputs
        block = SpectralConv(
            3, 4, modes[:dim], n_layers=1, output_scaling_factor=2)
    
        x = torch.randn(2, 3, *(12, )*dim)
        res = block(x)
        assert res.shape[1] == 4 # Check out channels
        assert(list(res.shape[2:]) == [12*2]*dim)
//  [python]
def test_reduce_scatter(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        partial_spec = [Partial()]

        origin_tensor = torch.randn(3 * self.world_size, 3)
        partial_tensor = distribute_tensor(origin_tensor, device_mesh, partial_spec)
        partial_placement = cast(Partial, partial_tensor.placements[0])

        sharded_tensor = mesh_reduce_scatter(
            partial_tensor._local_tensor, device_mesh, partial_placement.reduce_op, scatter_dim=0, mesh_dim=0
        )
        chunk_size = origin_tensor.size()[0] // self.world_size

        rank = device_mesh.get_rank()
        torch.allclose(origin_tensor[rank * chunk_size : (rank + 1) * chunk_size, :], sharded_tensor.cpu())
//  [python]
def test_reduce_scatter(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        partial_spec = [Partial()]

        origin_tensor = torch.randn(3 * self.world_size, 3)
        partial_tensor = distribute_tensor(origin_tensor, device_mesh, partial_spec)
        partial_placement = cast(Partial, partial_tensor.placements[0])

        sharded_tensor = mesh_reduce_scatter(
            partial_tensor._local_tensor, device_mesh, partial_placement.reduce_op, scatter_dim=0, mesh_dim=0
        )
        chunk_size = origin_tensor.size()[0] // self.world_size

        rank = device_mesh.get_rank()
        torch.testing.assert_close(origin_tensor[rank * chunk_size : (rank + 1) * chunk_size, :], sharded_tensor.cpu())
//  [python]
def test_all_reduce(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        partial_spec = [Partial()]

        origin_tensor = torch.randn(3 * self.world_size, 3)
        partial_tensor = distribute_tensor(origin_tensor, device_mesh, partial_spec)
        partial_placement = cast(Partial, partial_tensor.placements[0])

        reduced_tensor = mesh_all_reduce(
            partial_tensor._local_tensor, device_mesh, partial_placement.reduce_op, mesh_dim=0
        )
        torch.testing.assert_allclose(origin_tensor, reduced_tensor.cpu())
//  [python]
def test_all_reduce(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        partial_spec = [Partial()]

        origin_tensor = torch.randn(3 * self.world_size, 3)
        partial_tensor = distribute_tensor(origin_tensor, device_mesh, partial_spec)
        partial_placement = cast(Partial, partial_tensor.placements[0])

        reduced_tensor = mesh_all_reduce(
            partial_tensor._local_tensor, device_mesh, partial_placement.reduce_op, mesh_dim=0
        )
        torch.testing.assert_close(origin_tensor, reduced_tensor.cpu())
//  [python]
def test_all_gather(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        shard_spec = [Shard(0)]

        origin_tensor = torch.rand(3 * self.world_size, 3)
        sharded_tensor = distribute_tensor(origin_tensor, device_mesh, shard_spec)
        gathered_tensor = mesh_all_gather(
            sharded_tensor._local_tensor, origin_tensor.size(), device_mesh, scatter_dim=0, mesh_dim=0
        )

        torch.testing.assert_allclose(origin_tensor, gathered_tensor.cpu())
//  [python]
def test_all_gather(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        shard_spec = [Shard(0)]

        origin_tensor = torch.rand(3 * self.world_size, 3)
        sharded_tensor = distribute_tensor(origin_tensor, device_mesh, shard_spec)
        gathered_tensor = mesh_all_gather(
            sharded_tensor._local_tensor, origin_tensor.size(), device_mesh, scatter_dim=0, mesh_dim=0
        )

        torch.testing.assert_close(origin_tensor, gathered_tensor.cpu())
//  [python]
def test_all_gather(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        origin_tensor = torch.randn(3 * self.world_size, 3)
        shard_spec = [Shard(0)]

        sharded_tensor = distribute_tensor(origin_tensor, device_mesh=device_mesh, placements=shard_spec)

        for async_op in [True, False]:
            gathered_tensor = vescale_all_gather(sharded_tensor, mesh_dims=0, async_op=async_op)
            if async_op:
                self.assertTrue(isinstance(gathered_tensor._local_tensor, funcol.AsyncCollectiveTensor))

            self.assertEqual(origin_tensor.shape, gathered_tensor._local_tensor.shape)
            self.assertTrue(torch.allclose(origin_tensor, gathered_tensor._local_tensor.cpu()))
//  [python]
def test_all_gather(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        origin_tensor = torch.randn(3 * self.world_size, 3)
        shard_spec = [Shard(0)]

        sharded_tensor = distribute_tensor(origin_tensor, device_mesh=device_mesh, placements=shard_spec)

        for async_op in [True, False]:
            gathered_tensor = vescale_all_gather(sharded_tensor, mesh_dims=0, async_op=async_op)
            if async_op:
                self.assertTrue(isinstance(gathered_tensor._local_tensor, funcol.AsyncCollectiveTensor))

            torch.testing.assert_close(origin_tensor, gathered_tensor._local_tensor.cpu())
//  [python]
def test_autocast_fused_layer_norm(self, dtype, elementwise_affine, memory_efficient):
        bf16_fwd_thresholds = dict(rtol=1.6e-2, atol=3e-4)
        bf16_bwd_thresholds = dict(rtol=1.6e-2, atol=3e-3)
        batch_size = 16
        normalized_shape = [32, 16]
        native = torch.nn.LayerNorm(
            normalized_shape=normalized_shape, elementwise_affine=elementwise_affine
        ).to(device="cuda", dtype=dtype)
        fused = FusedLayerNorm(
            normalized_shape=normalized_shape, elementwise_affine=elementwise_affine, memory_efficient=memory_efficient
        ).cuda()
        native_x, fused_x = _prep_inputs(batch_size, normalized_shape, dtype)

        expected = native(native_x)
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = {'rtol': None, 'atol': None} if dtype == torch.half else bf16_fwd_thresholds
        torch.testing.assert_allclose(actual, expected, **tols, check_dtype=False) 

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.clone()
        expected.backward(g_native)
        actual.backward(g_fused)

        if dtype != torch.half:
            tols = bf16_bwd_thresholds
        elif memory_efficient:
            tols = {'rtol': 1e-3, 'atol': 1e-4}
        else:
            tols = {'rtol': None, 'atol': None}
        torch.testing.assert_allclose(native_x.grad, fused_x.grad, **tols, check_dtype=False)
//  [python]
def test_autocast_fused_layer_norm(self, dtype, elementwise_affine, memory_efficient):
        bf16_fwd_thresholds = dict(rtol=1.6e-2, atol=3e-4)
        bf16_bwd_thresholds = dict(rtol=1.6e-2, atol=3e-3)
        batch_size = 16
        normalized_shape = [32, 16]
        native = torch.nn.LayerNorm(
            normalized_shape=normalized_shape, elementwise_affine=elementwise_affine
        ).to(device="cuda", dtype=dtype)
        fused = FusedLayerNorm(
            normalized_shape=normalized_shape, elementwise_affine=elementwise_affine, memory_efficient=memory_efficient
        ).cuda()
        native_x, fused_x = _prep_inputs(batch_size, normalized_shape, dtype)

        expected = native(native_x)
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = {'rtol': None, 'atol': None} if dtype == torch.half else bf16_fwd_thresholds
        # original tests used torch.testing.assert_allclose, which disables dtype checking by default. 
        # link to issue here: https://github.com/pytorch/pytorch/issues/61844
        torch.testing.assert_close(actual, expected, **tols, check_dtype=False) 

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.clone()
        expected.backward(g_native)
        actual.backward(g_fused)

        if dtype != torch.half:
            tols = bf16_bwd_thresholds
        elif memory_efficient:
            tols = {'rtol': 1e-3, 'atol': 1e-4}
        else:
            tols = {'rtol': None, 'atol': None}
        torch.testing.assert_close(native_x.grad, fused_x.grad, **tols, check_dtype=False)
//  [python]
def test_autocast_fused_rms_norm(self, dtype, elementwise_affine, memory_efficient):
        bf16_fwd_thresholds = dict(rtol=1.6e-2, atol=3e-4)
        bf16_bwd_thresholds = dict(rtol=1.6e-2, atol=3e-3)
        batch_size = 16
        normalized_shape = [32, 16]
        native = FusedRMSNorm(
            normalized_shape=normalized_shape, elementwise_affine=elementwise_affine, memory_efficient=memory_efficient, 
        ).to(dtype=dtype)
        fused = FusedRMSNorm(
            normalized_shape=normalized_shape, elementwise_affine=elementwise_affine, memory_efficient=memory_efficient, 
        ).cuda()
        native_x, fused_x = _prep_inputs(batch_size, normalized_shape, dtype)

        expected = native(native_x.cpu())
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = {'rtol': None, 'atol': None} if dtype == torch.half else bf16_fwd_thresholds
        torch.testing.assert_allclose(actual, expected.detach().clone().cuda(), **tols, check_dtype=False)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.detach().clone().cuda()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = {'rtol': 1e-3, 'atol': 1e-3} if dtype == torch.half else bf16_bwd_thresholds
        torch.testing.assert_allclose(native_x.grad.cuda(), fused_x.grad, **tols, check_dtype=False)
//  [python]
def test_autocast_fused_rms_norm(self, dtype, elementwise_affine, memory_efficient):
        bf16_fwd_thresholds = dict(rtol=1.6e-2, atol=3e-4)
        bf16_bwd_thresholds = dict(rtol=1.6e-2, atol=3e-3)
        batch_size = 16
        normalized_shape = [32, 16]
        native = FusedRMSNorm(
            normalized_shape=normalized_shape, elementwise_affine=elementwise_affine, memory_efficient=memory_efficient, 
        ).to(dtype=dtype)
        fused = FusedRMSNorm(
            normalized_shape=normalized_shape, elementwise_affine=elementwise_affine, memory_efficient=memory_efficient, 
        ).cuda()
        native_x, fused_x = _prep_inputs(batch_size, normalized_shape, dtype)

        expected = native(native_x.cpu())
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = {'rtol': None, 'atol': None} if dtype == torch.half else bf16_fwd_thresholds
        torch.testing.assert_close(actual, expected.detach().clone().cuda(), **tols, check_dtype=False)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.detach().clone().cuda()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = {'rtol': 1e-3, 'atol': 1e-3} if dtype == torch.half else bf16_bwd_thresholds
        torch.testing.assert_close(native_x.grad.cuda(), fused_x.grad, **tols, check_dtype=False)
//  [python]
def test_all_reduce(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        partial_spec = [Partial()]

        origin_tensor = torch.randn(3 * self.world_size, 3)
        partial_tensor = distribute_tensor(origin_tensor, device_mesh, partial_spec)

        for async_op in [True, False]:
            reduced_tensor = vescale_all_reduce(partial_tensor, mesh_dims=0, async_op=async_op)
            if async_op:
                self.assertTrue(isinstance(reduced_tensor._local_tensor, funcol.AsyncCollectiveTensor))
            torch.testing.assert_allclose(origin_tensor, reduced_tensor._local_tensor.cpu())
//  [python]
def test_all_reduce(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        partial_spec = [Partial()]

        origin_tensor = torch.randn(3 * self.world_size, 3)
        partial_tensor = distribute_tensor(origin_tensor, device_mesh, partial_spec)

        for async_op in [True, False]:
            reduced_tensor = vescale_all_reduce(partial_tensor, mesh_dims=0, async_op=async_op)
            if async_op:
                self.assertTrue(isinstance(reduced_tensor._local_tensor, funcol.AsyncCollectiveTensor))
            torch.testing.assert_close(origin_tensor, reduced_tensor._local_tensor.cpu())
//  [python]
def test_surgery_torchscript_train(input: Any, alg_cls: Type[Algorithm]):
    """Tests torchscript model in train mode."""
    if alg_cls in (Factorize, GhostBatchNorm, StochasticDepth):
        pytest.xfail('Unsupported')

    alg_kwargs = get_alg_kwargs(alg_cls)
    model = get_alg_model(alg_cls)

    surgery_method = get_surgery_method(alg_cls)

    alg_kwargs = algo_kwargs.get(surgery_method, alg_kwargs)

    surgery_method(model, **alg_kwargs)

    scripted_func = torch.jit.script(model)
    scripted_func.train()  # type: ignore (third-party)
    model.train()
    torch.testing.assert_allclose(scripted_func(input), model(input))
//  [python]
def test_surgery_torchscript_train(input: Any, alg_cls: Type[Algorithm]):
    """Tests torchscript model in train mode."""
    if alg_cls in (Factorize, GhostBatchNorm, StochasticDepth):
        pytest.xfail('Unsupported')

    alg_kwargs = get_alg_kwargs(alg_cls)
    model = get_alg_model(alg_cls)

    surgery_method = get_surgery_method(alg_cls)

    alg_kwargs = algo_kwargs.get(surgery_method, alg_kwargs)

    surgery_method(model, **alg_kwargs)

    scripted_func = torch.jit.script(model)
    scripted_func.train()  # type: ignore (third-party)
    model.train()
    torch.testing.assert_close(scripted_func(input), model(input))
//  [python]
def test_reduce_scatter(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        partial_spec = [Partial()]

        origin_tensor = torch.randn(3 * self.world_size, 3)
        partial_tensor = distribute_tensor(origin_tensor, device_mesh, partial_spec)
        for async_op in [True, False]:
            sharded_tensor = vescale_reduce_scatter(
                partial_tensor, reduce_mesh_dims=0, scatter_dims=0, mesh_dims=0, async_op=async_op
            )
            if async_op:
                self.assertTrue(isinstance(sharded_tensor._local_tensor, funcol.AsyncCollectiveTensor))

            chunk_size = origin_tensor.size()[0] // self.world_size
            rank = device_mesh.get_rank()
            torch.testing.assert_allclose(
                origin_tensor[rank * chunk_size : (rank + 1) * chunk_size], sharded_tensor._local_tensor.cpu()
            )
        pass
//  [python]
def test_reduce_scatter(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        partial_spec = [Partial()]

        origin_tensor = torch.randn(3 * self.world_size, 3)
        partial_tensor = distribute_tensor(origin_tensor, device_mesh, partial_spec)
        for async_op in [True, False]:
            sharded_tensor = vescale_reduce_scatter(
                partial_tensor, reduce_mesh_dims=0, scatter_dims=0, mesh_dims=0, async_op=async_op
            )
            if async_op:
                self.assertTrue(isinstance(sharded_tensor._local_tensor, funcol.AsyncCollectiveTensor))

            chunk_size = origin_tensor.size()[0] // self.world_size
            rank = device_mesh.get_rank()
            torch.testing.assert_close(
                origin_tensor[rank * chunk_size : (rank + 1) * chunk_size], sharded_tensor._local_tensor.cpu()
            )
        pass
//  [python]
def test_surgery_torchscript_eval(input: Any, alg_cls: Type[Algorithm]):
    """Tests torchscript model in eval mode."""
    if alg_cls is Factorize:
        pytest.xfail('Unsupported')

    surgery_method = get_surgery_method(alg_cls)

    alg_kwargs = get_alg_kwargs(alg_cls)
    model = get_alg_model(alg_cls)
    alg_kwargs = algo_kwargs.get(surgery_method, alg_kwargs)

    surgery_method(model, **alg_kwargs)

    scripted_func = torch.jit.script(model)
    scripted_func.eval()  # type: ignore (third-party)
    model.eval()
    t.assert_close(scripted_func(input), model(input))
//  [python]
def test_surgery_torchscript_eval(input: Any, alg_cls: Type[Algorithm]):
    """Tests torchscript model in eval mode."""
    if alg_cls is Factorize:
        pytest.xfail('Unsupported')

    surgery_method = get_surgery_method(alg_cls)

    alg_kwargs = get_alg_kwargs(alg_cls)
    model = get_alg_model(alg_cls)
    alg_kwargs = algo_kwargs.get(surgery_method, alg_kwargs)

    surgery_method(model, **alg_kwargs)

    scripted_func = torch.jit.script(model)
    scripted_func.eval()  # type: ignore (third-party)
    model.eval()
    torch.testing.assert_close(scripted_func(input), model(input))
//  [python]
output = compiled_model(x)
expected_output = model(x1)

print(output)
print(expected_output)
torch.testing.assert_allclose(output, expected_output)
//  [python]
output = compiled_model(x)
expected_output = model(x1)

print(output)
print(expected_output)
torch.testing.assert_close(output, expected_output)
//  [python]
def test_surgery_torchfx_eval(
    input: Any,
    alg_cls: Type[Algorithm],
):
    """Tests torch.fx model in eval mode."""

    alg_kwargs = get_alg_kwargs(alg_cls)
    model = get_alg_model(alg_cls)
    surgery_method = get_surgery_method(alg_cls)

    if alg_cls in (BlurPool, GhostBatchNorm):
        pytest.xfail('Control flow')

    alg_kwargs = algo_kwargs.get(surgery_method, alg_kwargs)

    surgery_method(model, **alg_kwargs)

    model.eval()

    traced_func = torch.fx.symbolic_trace(model)
    torch.testing.assert_allclose(traced_func(input), model(input))
//  [python]
def test_surgery_torchfx_eval(
    input: Any,
    alg_cls: Type[Algorithm],
):
    """Tests torch.fx model in eval mode."""

    alg_kwargs = get_alg_kwargs(alg_cls)
    model = get_alg_model(alg_cls)
    surgery_method = get_surgery_method(alg_cls)

    if alg_cls in (BlurPool, GhostBatchNorm):
        pytest.xfail('Control flow')

    alg_kwargs = algo_kwargs.get(surgery_method, alg_kwargs)

    surgery_method(model, **alg_kwargs)

    model.eval()

    traced_func = torch.fx.symbolic_trace(model)
    torch.testing.assert_close(traced_func(input), model(input))
//  [python]
def assertAllClose(
        self,
        input: torch.Tensor,
        other: Union[torch.Tensor, float],
        rtol: float = 1e-05,
        atol: float = 1e-08,
        equal_nan: bool = False,
    ) -> None:
        r"""
        Calls torch.testing.assert_close, using the signature and default behavior
        of torch.allclose.

        Example output:
            AssertionError: Scalars are not close!

            Absolute difference: 1.0000034868717194 (up to 0.0001 allowed)
            Relative difference: 0.8348668001940709 (up to 1e-05 allowed)
        """
        # Why not just use the signature and behavior of `torch.testing.assert_close`?
        # Because we used `torch.allclose` for testing in the past, and the two don't
        # behave exactly the same. In particular, `assert_close` requires both `atol`
        # and `rtol` to be set if either one is.
        t.assert_close(
            input,
            other,
            rtol=rtol,
            atol=atol,
            equal_nan=equal_nan,
        )
//  [python]
def assertAllClose(
        self,
        input: torch.Tensor,
        other: Union[torch.Tensor, float],
        rtol: float = 1e-05,
        atol: float = 1e-08,
        equal_nan: bool = False,
    ) -> None:
        r"""
        Calls torch.testing.assert_close, using the signature and default behavior
        of torch.allclose.

        Example output:
            AssertionError: Scalars are not close!

            Absolute difference: 1.0000034868717194 (up to 0.0001 allowed)
            Relative difference: 0.8348668001940709 (up to 1e-05 allowed)
        """
        # Why not just use the signature and behavior of `torch.testing.assert_close`?
        # Because we used `torch.allclose` for testing in the past, and the two don't
        # behave exactly the same. In particular, `assert_close` requires both `atol`
        # and `rtol` to be set if either one is.
        torch.testing.assert_close(
            input,
            other,
            rtol=rtol,
            atol=atol,
            equal_nan=equal_nan,
        )
//  [python]
def test_add_bias_interleave_int4(self):
        # packed repr for -7 to 8
        packed_int4s = torch.tensor([[-104, -70, -36, -2, 16, 50, 84, 118]], dtype=torch.int8)
        actual_processed_int4s = self.add_bias_and_interleave_int4s(packed_int4s)
        # Packed repr for preprocessed cuda input (computed by hand)
        expected_processed_int4 = torch.tensor([[32, 100, 49, 117, -88, -20, -71, -3]], dtype=torch.int8)
        torch.testing.assert_allclose(actual_processed_int4s, expected_processed_int4, rtol=0, atol=0)
//  [python]
def test_add_bias_interleave_int4(self):
        # packed repr for -7 to 8
        packed_int4s = torch.tensor([[-104, -70, -36, -2, 16, 50, 84, 118]], dtype=torch.int8)
        actual_processed_int4s = self.add_bias_and_interleave_int4s(packed_int4s)
        # Packed repr for preprocessed cuda input (computed by hand)
        expected_processed_int4 = torch.tensor([[32, 100, 49, 117, -88, -20, -71, -3]], dtype=torch.int8)
        torch.testing.assert_close(actual_processed_int4s, expected_processed_int4, rtol=0, atol=0)
//  [python]
def test_surgery_onnx(
    input: Any,
    alg_cls: Type[Algorithm],
    tmp_path: pathlib.Path,
):
    """Tests onnx export and runtime"""
    pytest.importorskip('onnx')
    pytest.importorskip('onnxruntime')
    import onnx
    import onnxruntime as ort

    surgery_method = get_surgery_method(alg_cls)

    model = get_alg_model(alg_cls)
    alg_kwargs = get_alg_kwargs(alg_cls)
    alg_kwargs = algo_kwargs.get(surgery_method, alg_kwargs)

    surgery_method(model, **alg_kwargs)
    model.eval()

    onnx_path = os.path.join(tmp_path, 'model.onnx')
    torch.onnx.export(
        model,
        (input,),
        onnx_path,
        input_names=['input'],
        output_names=['output'],
    )

    # check onnx model
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)  # type: ignore (third-party)

    # run inference
    ort_session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
    outputs = ort_session.run(
        None,
        {'input': input[0].numpy()},
    )

    t.assert_allclose(
        outputs[0],
        model(input).detach().numpy(),
        rtol=1e-4,  # lower tolerance for ONNX
        atol=1e-3,  # lower tolerance for ONNX
    )
//  [python]
def test_surgery_onnx(
    input: Any,
    alg_cls: Type[Algorithm],
    tmp_path: pathlib.Path,
):
    """Tests onnx export and runtime"""
    pytest.importorskip('onnx')
    pytest.importorskip('onnxruntime')
    import onnx
    import onnxruntime as ort

    surgery_method = get_surgery_method(alg_cls)

    model = get_alg_model(alg_cls)
    alg_kwargs = get_alg_kwargs(alg_cls)
    alg_kwargs = algo_kwargs.get(surgery_method, alg_kwargs)

    surgery_method(model, **alg_kwargs)
    model.eval()

    onnx_path = os.path.join(tmp_path, 'model.onnx')
    torch.onnx.export(
        model,
        (input,),
        onnx_path,
        input_names=['input'],
        output_names=['output'],
    )

    # check onnx model
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)  # type: ignore (third-party)

    # run inference
    ort_session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
    outputs = ort_session.run(
        None,
        {'input': input[0].numpy()},
    )

    torch.testing.assert_close(
        outputs[0],
        model(input).detach().numpy(),
        rtol=1e-4,  # lower tolerance for ONNX
        atol=1e-3,  # lower tolerance for ONNX
    )
//  [python]
def test_unpack(self):
        packed_int4s = torch.tensor([[-104, -70, -36, -2, 16, 50, 84, 118]], dtype=torch.int8)

        unpacked_int4s_as_int8 = self.unpack_packed_int4s(packed_int4s)
        expected = torch.arange(-8, 8, dtype=torch.int8).reshape([1, 16])
        torch.testing.assert_allclose(unpacked_int4s_as_int8, expected, rtol=0, atol=0)
//  [python]
def test_unpack(self):
        packed_int4s = torch.tensor([[-104, -70, -36, -2, 16, 50, 84, 118]], dtype=torch.int8)

        unpacked_int4s_as_int8 = self.unpack_packed_int4s(packed_int4s)
        expected = torch.arange(-8, 8, dtype=torch.int8).reshape([1, 16])
        torch.testing.assert_close(unpacked_int4s_as_int8, expected, rtol=0, atol=0)
//  [python]
def test_add_bias_interleave_int8(self):
        int8s = torch.tensor([[-104, -70, -36, 127, 16, 50, 84, 118]], dtype=torch.int8)
        actual_processed_int8s = self.add_bias_and_interleave_int8s(int8s)
        # Packed repr for preprocessed cuda input (computed by hand)
        tmp = torch.tensor([[-104, -36, -70, 127, 16, 84, 50, 118]], dtype=torch.int32) + 128
        expected_processed_int8 = tmp.to(torch.int8)
        torch.testing.assert_allclose(actual_processed_int8s, expected_processed_int8, rtol=0, atol=0)
//  [python]
def test_add_bias_interleave_int8(self):
        int8s = torch.tensor([[-104, -70, -36, 127, 16, 50, 84, 118]], dtype=torch.int8)
        actual_processed_int8s = self.add_bias_and_interleave_int8s(int8s)
        # Packed repr for preprocessed cuda input (computed by hand)
        tmp = torch.tensor([[-104, -36, -70, 127, 16, 84, 50, 118]], dtype=torch.int32) + 128
        expected_processed_int8 = tmp.to(torch.int8)
        torch.testing.assert_close(actual_processed_int8s, expected_processed_int8, rtol=0, atol=0)
//  [python]
def interleave_tester(self, quant_type, arch):
        assert quant_type == torch.int8 or quant_type == torch.quint4x2
        experts  = [1, 4, 8]
        m_shapes = [128, 256, 1024]
        n_shapes = [128, 512, 1024, 4096]

        lower_bound = -128 if quant_type == torch.int8 else -8
        upper_bound = 127 if quant_type == torch.int8 else 7

        ref_impl = self.reference_interleave

        for expert in experts:
            for m_shape in m_shapes:
                for n_shape in n_shapes:
                    tensor = torch.randint(lower_bound, upper_bound, [expert, m_shape, n_shape], dtype=torch.int8)
                    ref_result = ref_impl(tensor.reshape([expert * m_shape, n_shape]), quant_type).reshape([expert, m_shape, n_shape])

                    if quant_type == torch.quint4x2:
                        tensor = self.pack_int4s(tensor)

                    act_result = self.quant_interleave(tensor, quant_type, arch)

                    if quant_type == torch.quint4x2:
                        act_result = self.unpack_packed_int4s(act_result)

                    torch.testing.assert_allclose(act_result, ref_result, rtol=0, atol=0)
//  [python]
def interleave_tester(self, quant_type, arch):
        assert quant_type == torch.int8 or quant_type == torch.quint4x2
        experts  = [1, 4, 8]
        m_shapes = [128, 256, 1024]
        n_shapes = [128, 512, 1024, 4096]

        lower_bound = -128 if quant_type == torch.int8 else -8
        upper_bound = 127 if quant_type == torch.int8 else 7

        ref_impl = self.reference_interleave

        for expert in experts:
            for m_shape in m_shapes:
                for n_shape in n_shapes:
                    tensor = torch.randint(lower_bound, upper_bound, [expert, m_shape, n_shape], dtype=torch.int8)
                    ref_result = ref_impl(tensor.reshape([expert * m_shape, n_shape]), quant_type).reshape([expert, m_shape, n_shape])

                    if quant_type == torch.quint4x2:
                        tensor = self.pack_int4s(tensor)

                    act_result = self.quant_interleave(tensor, quant_type, arch)

                    if quant_type == torch.quint4x2:
                        act_result = self.unpack_packed_int4s(act_result)

                    torch.testing.assert_close(act_result, ref_result, rtol=0, atol=0)
//  [python]
def transpose_test_helper(self, quant_type):
        assert quant_type == torch.int8 or quant_type == torch.quint4x2
        experts  = [1, 4, 8]
        m_shapes = [128, 256, 1024]
        n_shapes = [128, 4096]

        lower_bound = -128 if quant_type == torch.int8 else -8
        upper_bound = 127 if quant_type == torch.int8 else 7

        for expert in experts:
            for m_shape in m_shapes:
                for n_shape in n_shapes:
                    tensor = torch.randint(lower_bound, upper_bound, [expert, m_shape, n_shape], dtype=torch.int8)

                    # We want to move the data, but not change the shape. The actual impl just changes to col major.
                    ref_result = tensor.permute([0, 2, 1]).reshape([expert, m_shape, n_shape])

                    if quant_type == torch.quint4x2:
                        tensor = self.pack_int4s(tensor)

                    act_result = self.subbyte_transpose(tensor, quant_type)

                    if quant_type == torch.quint4x2:
                        act_result = self.unpack_packed_int4s(act_result)

                    self.assertEqual(act_result, ref_result)
//  [python]
def transpose_test_helper(self, quant_type):
        assert quant_type == torch.int8 or quant_type == torch.quint4x2
        experts  = [1, 4, 8]
        m_shapes = [128, 256, 1024]
        n_shapes = [128, 4096]

        lower_bound = -128 if quant_type == torch.int8 else -8
        upper_bound = 127 if quant_type == torch.int8 else 7

        for expert in experts:
            for m_shape in m_shapes:
                for n_shape in n_shapes:
                    tensor = torch.randint(lower_bound, upper_bound, [expert, m_shape, n_shape], dtype=torch.int8)

                    # We want to move the data, but not change the shape. The actual impl just changes to col major.
                    ref_result = tensor.permute([0, 2, 1]).reshape([expert, m_shape, n_shape])

                    if quant_type == torch.quint4x2:
                        tensor = self.pack_int4s(tensor)

                    act_result = self.subbyte_transpose(tensor, quant_type)

                    if quant_type == torch.quint4x2:
                        act_result = self.unpack_packed_int4s(act_result)

                    torch.testing.assert_close(act_result, ref_result, rtol=0, atol=0)
//  [python]
def test_pack_unpack_identity(self):
        initial_vals = torch.randint(-8, 7, [128, 128], dtype=torch.int8)
        expected = self.unpack_packed_int4s(self.pack_int4s(initial_vals))
        self.assertEqual(initial_vals, expected)
//  [python]
def test_pack_unpack_identity(self):
        initial_vals = torch.randint(-8, 7, [128, 128], dtype=torch.int8)
        expected = self.unpack_packed_int4s(self.pack_int4s(initial_vals))
        torch.testing.assert_close(initial_vals, expected, rtol=0, atol=0)
//  [python]
def test_pack(self):
        unpacked_i4s = torch.arange(-8, 8, dtype=torch.int8).reshape([1, 16])
        packed_i4s = self.pack_int4s(unpacked_i4s)
        expected = torch.tensor([[-104, -70, -36, -2, 16, 50, 84, 118]], dtype=torch.int8)
        torch.testing.assert_allclose(packed_i4s, expected, rtol=0, atol=0)
//  [python]
def test_pack(self):
        unpacked_i4s = torch.arange(-8, 8, dtype=torch.int8).reshape([1, 16])
        packed_i4s = self.pack_int4s(unpacked_i4s)
        expected = torch.tensor([[-104, -70, -36, -2, 16, 50, 84, 118]], dtype=torch.int8)
        torch.testing.assert_close(packed_i4s, expected, rtol=0, atol=0)
//  [python]
def test_inv_remap_state_dict_hf_llama(model_name):
    config = llama_config_to_gpt2_config(
        AutoConfig.from_pretrained(model_name, trust_remote_code=True)
    )
    state_dict = state_dict_from_pretrained(model_name)
    # inv_remap_state_dict_hf_llama should be the inverse of remap_state_dict_hf_llama
    state_dict = {key: val for key, val in state_dict.items() if "rotary_emb.inv_freq" not in key}
    pretrained_state_dict = remap_state_dict_hf_llama(state_dict, config)
    state_dict_recover = inv_remap_state_dict_hf_llama(pretrained_state_dict, config)
    assert set(state_dict_recover.keys()) == set(state_dict.keys())
    for key in state_dict_recover.keys():
        torch.testing.assert_allclose(state_dict_recover[key], state_dict[key])
//  [python]
def test_inv_remap_state_dict_hf_llama(model_name):
    config = llama_config_to_gpt2_config(
        AutoConfig.from_pretrained(model_name, trust_remote_code=True)
    )
    state_dict = state_dict_from_pretrained(model_name)
    # inv_remap_state_dict_hf_llama should be the inverse of remap_state_dict_hf_llama
    state_dict = {key: val for key, val in state_dict.items() if "rotary_emb.inv_freq" not in key}
    pretrained_state_dict = remap_state_dict_hf_llama(state_dict, config)
    state_dict_recover = inv_remap_state_dict_hf_llama(pretrained_state_dict, config)
    assert set(state_dict_recover.keys()) == set(state_dict.keys())
    for key in state_dict_recover.keys():
        torch.testing.assert_close(state_dict_recover[key], state_dict[key])
//  [python]
def test_lazy_load_basic(lit_llama):
    import lit_llama.utils

    with tempfile.TemporaryDirectory() as tmpdirname:
        m = torch.nn.Linear(5, 3)
        path = pathlib.Path(tmpdirname)
        fn = str(path / "test.pt")
        torch.save(m.state_dict(), fn)
        with lit_llama.utils.lazy_load(fn) as sd_lazy:
            assert "NotYetLoadedTensor" in str(next(iter(sd_lazy.values())))
            m2 = torch.nn.Linear(5, 3)
            m2.load_state_dict(sd_lazy)

        x = torch.randn(2, 5)
        actual = m2(x)
        expected = m(x)
        t.assert_close(actual, expected)
//  [python]
def test_lazy_load_basic(lit_llama):
    import lit_llama.utils

    with tempfile.TemporaryDirectory() as tmpdirname:
        m = torch.nn.Linear(5, 3)
        path = pathlib.Path(tmpdirname)
        fn = str(path / "test.pt")
        torch.save(m.state_dict(), fn)
        with lit_llama.utils.lazy_load(fn) as sd_lazy:
            assert "NotYetLoadedTensor" in str(next(iter(sd_lazy.values())))
            m2 = torch.nn.Linear(5, 3)
            m2.load_state_dict(sd_lazy)

        x = torch.randn(2, 5)
        actual = m2(x)
        expected = m(x)
        torch.testing.assert_close(actual, expected)