jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_sdxl_latent_scales(latent_mean, latent_std):
    model = stable_diffusion_xl(pretrained=False,
                                mask_pad_tokens=True,
                                fsdp=False,
                                encode_latents_in_fp16=False,
                                use_xformers=False,
                                latent_mean=latent_mean,
                                latent_std=latent_std)
    if isinstance(latent_mean, float):
        latent_mean_comp = [latent_mean] * 4
    else:
        latent_mean_comp = latent_mean
    latent_mean_comp = torch.tensor(latent_mean_comp).view(1, -1, 1, 1).to(model.latent_mean.device)
    if isinstance(latent_std, float):
        latent_std_comp = [latent_std] * 4
    else:
        latent_std_comp = latent_std
    latent_std_comp = torch.tensor(latent_std_comp).view(1, -1, 1, 1).to(model.latent_std.device)

    torch.testing.assert_allclose(model.latent_mean, latent_mean_comp)
    torch.testing.assert_allclose(model.latent_std, latent_std_comp)
//  [python]
def test_sdxl_latent_scales(latent_mean, latent_std):
    model = stable_diffusion_xl(pretrained=False,
                                mask_pad_tokens=True,
                                fsdp=False,
                                encode_latents_in_fp16=False,
                                use_xformers=False,
                                latent_mean=latent_mean,
                                latent_std=latent_std)
    if isinstance(latent_mean, float):
        latent_mean_comp = [latent_mean] * 4
    else:
        latent_mean_comp = latent_mean
    latent_mean_comp = torch.tensor(latent_mean_comp).view(1, -1, 1, 1).to(model.latent_mean.device)
    if isinstance(latent_std, float):
        latent_std_comp = [latent_std] * 4
    else:
        latent_std_comp = latent_std
    latent_std_comp = torch.tensor(latent_std_comp).view(1, -1, 1, 1).to(model.latent_std.device)

    torch.testing.assert_close(model.latent_mean, latent_mean_comp)
    torch.testing.assert_close(model.latent_std, latent_std_comp)
//  [python]
x, y = torch.randn(5, 3), torch.randn(5, 3)
orig_output = traced(x, y)

wrap_in_activation_function(traced, ActivationFunction.LEAKY_RELU)
new_output = traced(x, y)

torch.testing.assert_allclose(new_output, torch.nn.LeakyReLU()(orig_output))
//  [python]
x, y = torch.randn(5, 3), torch.randn(5, 3)
orig_output = traced(x, y)

wrap_in_activation_function(traced, ActivationFunction.LEAKY_RELU)
new_output = traced(x, y)

torch.testing.assert_close(new_output, torch.nn.LeakyReLU()(orig_output))
//  [python]
def test_spherical_harmonics(components):
    torch.manual_seed(0)
    N = 1000000

    dx = torch.normal(0, 1, size=(N, 3))
    dx = dx / torch.linalg.norm(dx, dim=-1, keepdim=True)
    sh = components_from_spherical_harmonics(components, dx)
    matrix = (sh.T @ sh) / N * 4 * torch.pi
    torch.testing.assert_allclose(matrix, torch.eye(components**2), rtol=0, atol=1.5e-2)
//  [python]
def test_spherical_harmonics(components):
    torch.manual_seed(0)
    N = 1000000

    dx = torch.normal(0, 1, size=(N, 3))
    dx = dx / torch.linalg.norm(dx, dim=-1, keepdim=True)
    sh = components_from_spherical_harmonics(components, dx)
    matrix = (sh.T @ sh) / N * 4 * torch.pi
    torch.testing.assert_close(matrix, torch.eye(components**2), rtol=0, atol=1.5e-2)
//  [python]
def test_mode():
    """Test the mode either returns a lower triangle or a matrix of zeros."""
    p = torch.tensor(0.6)
    mode = ErdosRenyiDAGDistribution(num_nodes=5, probs=p).mode
    torch.testing.assert_allclose(mode, torch.tril(torch.ones_like(mode), diagonal=-1))

    p = torch.tensor(0.2)
    mode = ErdosRenyiDAGDistribution(num_nodes=5, probs=p).mode
    torch.testing.assert_allclose(mode, torch.zeros_like(mode))
//  [python]
def test_mode():
    """Test the mode either returns a lower triangle or a matrix of zeros."""
    p = torch.tensor(0.6)
    mode = ErdosRenyiDAGDistribution(num_nodes=5, probs=p).mode
    torch.testing.assert_close(mode, torch.tril(torch.ones_like(mode), diagonal=-1))

    p = torch.tensor(0.2)
    mode = ErdosRenyiDAGDistribution(num_nodes=5, probs=p).mode
    torch.testing.assert_close(mode, torch.zeros_like(mode))
//  [python]
def test_extreme_sample():
    """Test that extreme probabilities give rise to expected graphs"""
    n = 6
    p = torch.tensor(1.0)
    sample = ErdosRenyiDAGDistribution(num_nodes=n, probs=p).sample()
    t.assert_close(sample.sum(dim=(-2, -1)).item(), n * (n - 1) / 2.0)

    p = torch.tensor(0.0)
    sample = ErdosRenyiDAGDistribution(num_nodes=n, probs=p).sample()
    t.assert_close(sample, torch.zeros_like(sample))
//  [python]
def test_extreme_sample():
    """Test that extreme probabilities give rise to expected graphs"""
    n = 6
    p = torch.tensor(1.0)
    sample = ErdosRenyiDAGDistribution(num_nodes=n, probs=p).sample()
    torch.testing.assert_close(sample.sum(dim=(-2, -1)).item(), n * (n - 1) / 2.0)

    p = torch.tensor(0.0)
    sample = ErdosRenyiDAGDistribution(num_nodes=n, probs=p).sample()
    torch.testing.assert_close(sample, torch.zeros_like(sample))
//  [python]
def test_simple_sum(self):
        n = 32
        cg = construct_adder(n)

        tA = torch.randn(n)
        tB = torch.randn(n)
        tC = torch.empty(n)
        cg.call([tA, tB, tC])
        t.assert_allclose(tA + tB, tC)
//  [python]
def test_simple_sum(self):
        n = 32
        cg = construct_adder(n)

        tA = torch.randn(n)
        tB = torch.randn(n)
        tC = torch.empty(n)
        cg.call([tA, tB, tC])
        torch.testing.assert_close(tA + tB, tC)
//  [python]
def test_asynchronous_complete_cumsum_2d(
        self,
        n: int,
        b: int,
        index_types: Tuple[Type[object], Type[object]],
        device: torch.device,
    ) -> None:
        (pt_index_dtype, np_index_dtype) = index_types

        # The CPU variants of asynchronous_*_cumsum support floats, since some
        # downstream tests appear to be relying on this behavior.  As such, the
        # test is disabled for GPU + float test cases.
        if device == torch.device("cuda") and pt_index_dtype is torch.float32:
            return

        # pyre-ignore-errors[16]
        x = torch.randint(low=0, high=100, size=(b, n)).type(pt_index_dtype).to(device)

        zc = torch.ops.fbgemm.asynchronous_complete_cumsum(x)
        zeros = torch.zeros(b, 1)
        t.assert_allclose(
            torch.from_numpy(
                np.cumsum(torch.concat([zeros, x.cpu()], dim=1).numpy(), axis=1).astype(
                    np_index_dtype
                )
            ),
            zc.cpu(),
        )
//  [python]
def test_asynchronous_complete_cumsum_2d(
        self,
        n: int,
        b: int,
        index_types: Tuple[Type[object], Type[object]],
        device: torch.device,
    ) -> None:
        (pt_index_dtype, np_index_dtype) = index_types

        # The CPU variants of asynchronous_*_cumsum support floats, since some
        # downstream tests appear to be relying on this behavior.  As such, the
        # test is disabled for GPU + float test cases.
        if device == torch.device("cuda") and pt_index_dtype is torch.float32:
            return

        # pyre-ignore-errors[16]
        x = torch.randint(low=0, high=100, size=(b, n)).type(pt_index_dtype).to(device)

        zc = torch.ops.fbgemm.asynchronous_complete_cumsum(x)
        zeros = torch.zeros(b, 1)
        torch.testing.assert_close(
            torch.from_numpy(
                np.cumsum(torch.concat([zeros, x.cpu()], dim=1).numpy(), axis=1).astype(
                    np_index_dtype
                )
            ),
            zc.cpu(),
        )
//  [python]
def test_num_deges():
    num_edges = 16
    samples = ErdosRenyiDAGDistribution(num_nodes=8, num_edges=torch.tensor(num_edges)).sample(torch.Size([100]))

    assert samples.shape == torch.Size([100, 8, 8])
    torch.testing.assert_allclose(
        samples.sum(dim=(-2, -1)).mean(), torch.tensor(num_edges, dtype=torch.float32), atol=2.0, rtol=0.1
    )

    samples = ErdosRenyiDAGDistribution(num_nodes=2, num_edges=torch.tensor(2)).sample(torch.Size([100]))

    assert samples.shape == torch.Size([100, 2, 2])
    torch.testing.assert_allclose(samples.sum(dim=(-2, -1)).mean(), torch.tensor(1, dtype=torch.float32))
//  [python]
def test_num_deges():
    num_edges = 16
    samples = ErdosRenyiDAGDistribution(num_nodes=8, num_edges=torch.tensor(num_edges)).sample(torch.Size([100]))

    assert samples.shape == torch.Size([100, 8, 8])
    torch.testing.assert_close(
        samples.sum(dim=(-2, -1)).mean(), torch.tensor(num_edges, dtype=torch.float32), atol=2.0, rtol=0.1
    )

    samples = ErdosRenyiDAGDistribution(num_nodes=2, num_edges=torch.tensor(2)).sample(torch.Size([100]))

    assert samples.shape == torch.Size([100, 2, 2])
    torch.testing.assert_close(samples.sum(dim=(-2, -1)).mean(), torch.tensor(1, dtype=torch.float32))
//  [python]
def test_fsdp_train_save_load(tmp_path, manual_wrapping, precision):
    """Test FSDP training, saving and loading with different wrapping and precision settings."""
    trainer_cls = _TrainerManualWrapping if manual_wrapping else _Trainer
    fabric = Fabric(
        accelerator="cuda",
        strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
        devices=2,
        precision=precision,
    )
    fabric.launch()
    trainer = trainer_cls(fabric)
    trainer.run()

    checkpoint_path = fabric.broadcast(str(tmp_path / "fsdp-checkpoint"))

    params_before = deepcopy(list(trainer.model.parameters()))
    state = {"model": trainer.model, "optimizer": trainer.optimizer, "steps": 1}
    fabric.save(checkpoint_path, state)
    assert set(os.listdir(checkpoint_path)) == {"meta.pt", ".metadata", "__0_0.distcp", "__1_0.distcp"}

    # re-init all objects and resume
    fabric = Fabric(
        accelerator="cuda",
        strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
        devices=2,
        precision=precision,
    )
    fabric.launch()
    trainer = trainer_cls(fabric)
    trainer.run()

    # check correctness with loaded state
    state = {"model": trainer.model, "optimizer": trainer.optimizer, "steps": 0}
    metadata = fabric.load(checkpoint_path, state)
    for p0, p1 in zip(params_before, trainer.model.parameters()):
        torch.testing.assert_allclose(p0, p1, atol=0, rtol=0, equal_nan=True)

    # check user data in state reloaded
    assert state["steps"] == 1
    assert not metadata

    # attempt to load a key not in the metadata checkpoint
    state = {"model": trainer.model, "coconut": 11}
    with pytest.raises(KeyError, match="The requested state contains a key 'coconut' that does not exist"):
        fabric.load(checkpoint_path, state)

    # `strict=False` ignores the missing key
    state = {"model": trainer.model, "coconut": 11}
    fabric.load(checkpoint_path, state, strict=False)
    assert state["coconut"] == 11
//  [python]
def test_fsdp_train_save_load(tmp_path, manual_wrapping, precision):
    """Test FSDP training, saving and loading with different wrapping and precision settings."""
    trainer_cls = _TrainerManualWrapping if manual_wrapping else _Trainer
    fabric = Fabric(
        accelerator="cuda",
        strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
        devices=2,
        precision=precision,
    )
    fabric.launch()
    trainer = trainer_cls(fabric)
    trainer.run()

    checkpoint_path = fabric.broadcast(str(tmp_path / "fsdp-checkpoint"))

    params_before = deepcopy(list(trainer.model.parameters()))
    state = {"model": trainer.model, "optimizer": trainer.optimizer, "steps": 1}
    fabric.save(checkpoint_path, state)
    assert set(os.listdir(checkpoint_path)) == {"meta.pt", ".metadata", "__0_0.distcp", "__1_0.distcp"}

    # re-init all objects and resume
    fabric = Fabric(
        accelerator="cuda",
        strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
        devices=2,
        precision=precision,
    )
    fabric.launch()
    trainer = trainer_cls(fabric)
    trainer.run()

    # check correctness with loaded state
    state = {"model": trainer.model, "optimizer": trainer.optimizer, "steps": 0}
    metadata = fabric.load(checkpoint_path, state)
    for p0, p1 in zip(params_before, trainer.model.parameters()):
        torch.testing.assert_close(p0, p1, atol=0, rtol=0, equal_nan=True)

    # check user data in state reloaded
    assert state["steps"] == 1
    assert not metadata

    # attempt to load a key not in the metadata checkpoint
    state = {"model": trainer.model, "coconut": 11}
    with pytest.raises(KeyError, match="The requested state contains a key 'coconut' that does not exist"):
        fabric.load(checkpoint_path, state)

    # `strict=False` ignores the missing key
    state = {"model": trainer.model, "coconut": 11}
    fabric.load(checkpoint_path, state, strict=False)
    assert state["coconut"] == 11
//  [python]
def test_call_raw(self):
        n = 16
        cg = construct_adder(n, dtype=torch.float64)

        tA = torch.randn(n, dtype=torch.float64)
        tB = torch.randn(n, dtype=torch.float64)
        tC = torch.empty(n, dtype=torch.float64)
        cg.call_raw([tA.data_ptr(), tB.data_ptr(), tC.data_ptr()])
        torch.testing.assert_allclose(tA + tB, tC)
//  [python]
def test_call_raw(self):
        n = 16
        cg = construct_adder(n, dtype=torch.float64)

        tA = torch.randn(n, dtype=torch.float64)
        tB = torch.randn(n, dtype=torch.float64)
        tC = torch.empty(n, dtype=torch.float64)
        cg.call_raw([tA.data_ptr(), tB.data_ptr(), tC.data_ptr()])
        torch.testing.assert_close(tA + tB, tC)
//  [python]
def test_cumsum(
        self,
        n: int,
        index_types: Tuple[Type[object], Type[object]],
        device: torch.device,
    ) -> None:
        (pt_index_dtype, np_index_dtype) = index_types

        # The CPU variants of asynchronous_*_cumsum support floats, since some
        # downstream tests appear to be relying on this behavior.  As such, the
        # test is disabled for GPU + float test cases.
        if device == torch.device("cuda") and pt_index_dtype is torch.float32:
            return

        # pyre-ignore-errors[16]
        x = torch.randint(low=0, high=100, size=(n,)).type(pt_index_dtype).to(device)
        ze = torch.ops.fbgemm.asynchronous_exclusive_cumsum(x)
        zi = torch.ops.fbgemm.asynchronous_inclusive_cumsum(x)
        zc = torch.ops.fbgemm.asynchronous_complete_cumsum(x)

        t.assert_close(
            torch.from_numpy(np.cumsum(x.cpu().numpy()).astype(np_index_dtype)),
            zi.cpu(),
        )
        t.assert_close(
            torch.from_numpy(
                (np.cumsum([0] + x.cpu().numpy().tolist())[:-1]).astype(np_index_dtype)
            ),
            ze.cpu(),
        )
        t.assert_close(
            torch.from_numpy(
                (np.cumsum([0] + x.cpu().numpy().tolist())).astype(np_index_dtype)
            ),
            zc.cpu(),
        )

        # meta tests
        # pyre-ignore-errors[16]
        mx = torch.randint(low=0, high=100, size=(n,)).type(pt_index_dtype).to("meta")

        mze = torch.ops.fbgemm.asynchronous_exclusive_cumsum(mx)
        self.assertEqual(ze.size(), mze.size())

        mzi = torch.ops.fbgemm.asynchronous_inclusive_cumsum(mx)
        self.assertEqual(zi.size(), mzi.size())

        mzc = torch.ops.fbgemm.asynchronous_complete_cumsum(mx)
        self.assertEqual(zc.size(), mzc.size())
//  [python]
def test_cumsum(
        self,
        n: int,
        index_types: Tuple[Type[object], Type[object]],
        device: torch.device,
    ) -> None:
        (pt_index_dtype, np_index_dtype) = index_types

        # The CPU variants of asynchronous_*_cumsum support floats, since some
        # downstream tests appear to be relying on this behavior.  As such, the
        # test is disabled for GPU + float test cases.
        if device == torch.device("cuda") and pt_index_dtype is torch.float32:
            return

        # pyre-ignore-errors[16]
        x = torch.randint(low=0, high=100, size=(n,)).type(pt_index_dtype).to(device)
        ze = torch.ops.fbgemm.asynchronous_exclusive_cumsum(x)
        zi = torch.ops.fbgemm.asynchronous_inclusive_cumsum(x)
        zc = torch.ops.fbgemm.asynchronous_complete_cumsum(x)

        torch.testing.assert_close(
            torch.from_numpy(np.cumsum(x.cpu().numpy()).astype(np_index_dtype)),
            zi.cpu(),
        )
        torch.testing.assert_close(
            torch.from_numpy(
                (np.cumsum([0] + x.cpu().numpy().tolist())[:-1]).astype(np_index_dtype)
            ),
            ze.cpu(),
        )
        torch.testing.assert_close(
            torch.from_numpy(
                (np.cumsum([0] + x.cpu().numpy().tolist())).astype(np_index_dtype)
            ),
            zc.cpu(),
        )

        # meta tests
        # pyre-ignore-errors[16]
        mx = torch.randint(low=0, high=100, size=(n,)).type(pt_index_dtype).to("meta")

        mze = torch.ops.fbgemm.asynchronous_exclusive_cumsum(mx)
        self.assertEqual(ze.size(), mze.size())

        mzi = torch.ops.fbgemm.asynchronous_inclusive_cumsum(mx)
        self.assertEqual(zi.size(), mzi.size())

        mzc = torch.ops.fbgemm.asynchronous_complete_cumsum(mx)
        self.assertEqual(zc.size(), mzc.size())
//  [python]
def test_external_calls(self):
        dtype = torch.float32

        A = te.BufHandle("A", [1, 4], dtype)
        B = te.BufHandle("B", [4, 1], dtype)
        C = te.BufHandle("C", [1, 1], dtype)

        s = te.ExternalCall(C, "nnc_aten_matmul", [A, B], [])

        loopnest = te.LoopNest(s, [C])
        loopnest.prepare_for_codegen()
        codegen = te.construct_codegen("ir_eval", s, [A, B, C])

        tA = torch.ones(1, 4)
        tB = torch.ones(4, 1)
        tC = torch.empty(1, 1)
        codegen.call([tA, tB, tC])
        torch.testing.assert_allclose(torch.matmul(tA, tB), tC)
//  [python]
def test_external_calls(self):
        dtype = torch.float32

        A = te.BufHandle("A", [1, 4], dtype)
        B = te.BufHandle("B", [4, 1], dtype)
        C = te.BufHandle("C", [1, 1], dtype)

        s = te.ExternalCall(C, "nnc_aten_matmul", [A, B], [])

        loopnest = te.LoopNest(s, [C])
        loopnest.prepare_for_codegen()
        codegen = te.construct_codegen("ir_eval", s, [A, B, C])

        tA = torch.ones(1, 4)
        tB = torch.ones(4, 1)
        tC = torch.empty(1, 1)
        codegen.call([tA, tB, tC])
        torch.testing.assert_close(torch.matmul(tA, tB), tC)
//  [python]
def test_dynamic_shape(self):
        dN = te.VarHandle(torch.int32)
        A = te.BufHandle([dN], torch.float64)
        B = te.BufHandle([dN], torch.float64)

        def compute(i):
            return A.load(i) - B.load(i)

        C = te.Compute("C", [dN], compute)

        loopnest = te.LoopNest([C])
        loopnest.prepare_for_codegen()

        cg = te.construct_codegen("ir_eval", loopnest.simplify(), [A, B, C, dN])

        def test_with_shape(n):
            tA = torch.randn(n, dtype=torch.double)
            tB = torch.randn(n, dtype=torch.double)
            tC = torch.empty(n, dtype=torch.double)
            cg.call([tA, tB, tC, n])
            t.assert_allclose(tA - tB, tC)

        test_with_shape(8)
        test_with_shape(31)
//  [python]
def test_dynamic_shape(self):
        dN = te.VarHandle(torch.int32)
        A = te.BufHandle([dN], torch.float64)
        B = te.BufHandle([dN], torch.float64)

        def compute(i):
            return A.load(i) - B.load(i)

        C = te.Compute("C", [dN], compute)

        loopnest = te.LoopNest([C])
        loopnest.prepare_for_codegen()

        cg = te.construct_codegen("ir_eval", loopnest.simplify(), [A, B, C, dN])

        def test_with_shape(n):
            tA = torch.randn(n, dtype=torch.double)
            tB = torch.randn(n, dtype=torch.double)
            tC = torch.empty(n, dtype=torch.double)
            cg.call([tA, tB, tC, n])
            torch.testing.assert_close(tA - tB, tC)

        test_with_shape(8)
        test_with_shape(31)
//  [python]
def test_wasserstein1d_valid_input(self) -> None:
        # Checking with distribution values only
        metric = Wasserstein1D()
        x = torch.tensor([5, -5, -7, 9, -3])
        y = torch.tensor([9, -7, 5, -4, -2])
        metric.update(x, y)
        result = metric.compute()
        expected = torch.tensor([0.39999999999999997])
        t.assert_close(
            result,
            expected,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        # Checking with distribution and weight values
        metric = Wasserstein1D()
        x = torch.tensor([-13, -9, -19, 11, -18, -20, 8, 2, -8, -18])
        y = torch.tensor([9, 6, -5, -11, 9, -4, -13, -19, -14, 4])
        x_weights = torch.tensor([3, 3, 1, 2, 2, 3, 2, 2, 2, 3])
        y_weights = torch.tensor([2, 2, 1, 1, 2, 2, 1, 1, 1, 1])
        metric.update(x, y, x_weights, y_weights)
        result = metric.compute()
        expected = torch.tensor([8.149068322981368])
        t.assert_close(
            result,
            expected,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        # Checking with different distribution shapes
        metric = Wasserstein1D()
        x = torch.tensor([5, -5, -7, 9, -3])
        y = torch.tensor([9, -7, 5, -4, -2, 4, -1])
        metric.update(x, y)
        result = metric.compute()
        expected = torch.tensor([1.4571428571428569])
        t.assert_close(
            result,
            expected,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        # Checking with identical distributions
        metric = Wasserstein1D()
        x = torch.tensor([-13, -9, -19, 11, -18, -20, 8, 2, -8, -18])
        x_weights = torch.tensor([3, 3, 1, 2, 2, 3, 2, 2, 2, 3])
        metric.update(x, x, x_weights, x_weights)
        result = metric.compute()
        expected = torch.tensor([0.0])
        t.assert_close(
            result,
            expected,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )
//  [python]
def test_wasserstein1d_valid_input(self) -> None:
        # Checking with distribution values only
        metric = Wasserstein1D()
        x = torch.tensor([5, -5, -7, 9, -3])
        y = torch.tensor([9, -7, 5, -4, -2])
        metric.update(x, y)
        result = metric.compute()
        expected = torch.tensor([0.39999999999999997])
        torch.testing.assert_close(
            result,
            expected,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        # Checking with distribution and weight values
        metric = Wasserstein1D()
        x = torch.tensor([-13, -9, -19, 11, -18, -20, 8, 2, -8, -18])
        y = torch.tensor([9, 6, -5, -11, 9, -4, -13, -19, -14, 4])
        x_weights = torch.tensor([3, 3, 1, 2, 2, 3, 2, 2, 2, 3])
        y_weights = torch.tensor([2, 2, 1, 1, 2, 2, 1, 1, 1, 1])
        metric.update(x, y, x_weights, y_weights)
        result = metric.compute()
        expected = torch.tensor([8.149068322981368])
        torch.testing.assert_close(
            result,
            expected,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        # Checking with different distribution shapes
        metric = Wasserstein1D()
        x = torch.tensor([5, -5, -7, 9, -3])
        y = torch.tensor([9, -7, 5, -4, -2, 4, -1])
        metric.update(x, y)
        result = metric.compute()
        expected = torch.tensor([1.4571428571428569])
        torch.testing.assert_close(
            result,
            expected,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        # Checking with identical distributions
        metric = Wasserstein1D()
        x = torch.tensor([-13, -9, -19, 11, -18, -20, 8, 2, -8, -18])
        x_weights = torch.tensor([3, 3, 1, 2, 2, 3, 2, 2, 2, 3])
        metric.update(x, x, x_weights, x_weights)
        result = metric.compute()
        expected = torch.tensor([0.0])
        torch.testing.assert_close(
            result,
            expected,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )
//  [python]
def test_dynamic_shape_2d(self):
        dN = te.VarHandle(torch.int32)
        dM = te.VarHandle(torch.int32)
        A = te.BufHandle([dN, dM], torch.float64)
        B = te.BufHandle([dN, dM], torch.float64)

        def compute(i, j):
            return A.load([i, j]) - B.load([i, j])

        C = te.Compute("C", [dN, dM], compute)

        loopnest = te.LoopNest([C])
        loopnest.prepare_for_codegen()

        cg = te.construct_codegen("ir_eval", loopnest.simplify(), [A, B, C, dN, dM])

        def test_with_shape(n, m):
            tA = torch.randn(n, m, dtype=torch.double)
            tB = torch.randn(n, m, dtype=torch.double)
            tC = torch.empty(n, m, dtype=torch.double)
            cg.call([tA, tB, tC, n, m])
            t.assert_allclose(tA - tB, tC)

        test_with_shape(2, 4)
        test_with_shape(5, 3)
//  [python]
def test_dynamic_shape_2d(self):
        dN = te.VarHandle(torch.int32)
        dM = te.VarHandle(torch.int32)
        A = te.BufHandle([dN, dM], torch.float64)
        B = te.BufHandle([dN, dM], torch.float64)

        def compute(i, j):
            return A.load([i, j]) - B.load([i, j])

        C = te.Compute("C", [dN, dM], compute)

        loopnest = te.LoopNest([C])
        loopnest.prepare_for_codegen()

        cg = te.construct_codegen("ir_eval", loopnest.simplify(), [A, B, C, dN, dM])

        def test_with_shape(n, m):
            tA = torch.randn(n, m, dtype=torch.double)
            tB = torch.randn(n, m, dtype=torch.double)
            tC = torch.empty(n, m, dtype=torch.double)
            cg.call([tA, tB, tC, n, m])
            torch.testing.assert_close(tA - tB, tC)

        test_with_shape(2, 4)
        test_with_shape(5, 3)
//  [python]
def test_all_gather(self):
    dist_world_size = xu.getenv_as('WORLD_SIZE', int)
    devices_per_thread = xr.addressable_device_count()

    expected_world_size = dist_world_size * devices_per_thread

    rank = torch.tensor([dist.get_rank()],
                        dtype=torch.float32,
                        device=xm.xla_device())
    output = [rank.clone() for _ in range(expected_world_size)]
    dist.all_gather(output, rank)
    result = torch.concat(output)
    xm.mark_step()

    expected = torch.arange(0, expected_world_size, step=1, dtype=torch.float32)
    torch.testing.assert_allclose(result.cpu(), expected)
//  [python]
def test_all_gather(self):
    dist_world_size = xu.getenv_as('WORLD_SIZE', int)
    devices_per_thread = xr.addressable_device_count()

    expected_world_size = dist_world_size * devices_per_thread

    rank = torch.tensor([dist.get_rank()],
                        dtype=torch.float32,
                        device=xm.xla_device())
    output = [rank.clone() for _ in range(expected_world_size)]
    dist.all_gather(output, rank)
    result = torch.concat(output)
    xm.mark_step()

    expected = torch.arange(0, expected_world_size, step=1, dtype=torch.float32)
    torch.testing.assert_close(result.cpu(), expected)
//  [python]
def test_all_reduce(self):
    # The test is inspired by https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_reduce
    dist_world_size = xu.getenv_as('WORLD_SIZE', int)
    devices_per_thread = xr.addressable_device_count()
    world_size = dist_world_size * devices_per_thread

    # If world_size=2, then the `tensors` below will be [[1, 2], [3, 4]].
    # The `expected` will be [4, 6].
    tensors = [
        torch.arange(2, dtype=torch.int64) + 1 + 2 * r
        for r in range(world_size)
    ]
    expected = sum(tensors)

    xla_tensor = torch.arange(
        2, dtype=torch.int64, device=xm.xla_device()) + 1 + 2 * dist.get_rank()
    dist.all_reduce(xla_tensor, op=dist.ReduceOp.SUM)
    xm.mark_step()

    torch.testing.assert_allclose(xla_tensor.cpu(), expected)
//  [python]
def test_all_reduce(self):
    # The test is inspired by https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_reduce
    dist_world_size = xu.getenv_as('WORLD_SIZE', int)
    devices_per_thread = xr.addressable_device_count()
    world_size = dist_world_size * devices_per_thread

    # If world_size=2, then the `tensors` below will be [[1, 2], [3, 4]].
    # The `expected` will be [4, 6].
    tensors = [
        torch.arange(2, dtype=torch.int64) + 1 + 2 * r
        for r in range(world_size)
    ]
    expected = sum(tensors)

    xla_tensor = torch.arange(
        2, dtype=torch.int64, device=xm.xla_device()) + 1 + 2 * dist.get_rank()
    dist.all_reduce(xla_tensor, op=dist.ReduceOp.SUM)
    xm.mark_step()

    torch.testing.assert_close(xla_tensor.cpu(), expected)
//  [python]
def test_lower_bound():
    x = torch.rand(16, requires_grad=True)

    bound = torch.rand(1)

    t.assert_close(
        lower_bound(x, bound),
        torch.max(x, bound),
    )

    bound = torch.rand(1)

    y = lower_bound(x, bound)

    t.assert_close(
        y,
        torch.max(x, bound),
    )

    y.backward(x)

    assert x.grad is not None

    t.assert_close(
        x.grad,
        (x >= bound) * x,
    )
//  [python]
def test_lower_bound():
    x = torch.rand(16, requires_grad=True)

    bound = torch.rand(1)

    torch.testing.assert_close(
        lower_bound(x, bound),
        torch.max(x, bound),
    )

    bound = torch.rand(1)

    y = lower_bound(x, bound)

    torch.testing.assert_close(
        y,
        torch.max(x, bound),
    )

    y.backward(x)

    assert x.grad is not None

    torch.testing.assert_close(
        x.grad,
        (x >= bound) * x,
    )
//  [python]
def test_reduce_scatter(self):
    # The test is inspired by https://pytorch.org/docs/stable/distributed.html#torch.distributed.reduce_scatter
    dist_world_size = xu.getenv_as('WORLD_SIZE', int)
    devices_per_thread = xr.addressable_device_count()
    world_size = dist_world_size * devices_per_thread
    # If world_size=2, then `tensor` will be tensor([0, 2, 4, 6])
    # `expected` will be [0, 2] on rank 0 and [4, 6] on rank 1.
    tensor = world_size * torch.arange(
        world_size * world_size, dtype=torch.int64)
    expected = torch.split(tensor, world_size)[dist.get_rank()]

    tensor_out = torch.zeros(
        world_size, dtype=torch.int64, device=xm.xla_device())
    tensor_in = torch.arange(
        world_size * world_size, dtype=torch.int64, device=xm.xla_device())
    dist.reduce_scatter(tensor_out, [tensor_in], op=dist.ReduceOp.SUM)
    xm.mark_step()

    torch.testing.assert_allclose(tensor_out.cpu(), expected)
//  [python]
def test_reduce_scatter(self):
    # The test is inspired by https://pytorch.org/docs/stable/distributed.html#torch.distributed.reduce_scatter
    dist_world_size = xu.getenv_as('WORLD_SIZE', int)
    devices_per_thread = xr.addressable_device_count()
    world_size = dist_world_size * devices_per_thread
    # If world_size=2, then `tensor` will be tensor([0, 2, 4, 6])
    # `expected` will be [0, 2] on rank 0 and [4, 6] on rank 1.
    tensor = world_size * torch.arange(
        world_size * world_size, dtype=torch.int64)
    expected = torch.split(tensor, world_size)[dist.get_rank()]

    tensor_out = torch.zeros(
        world_size, dtype=torch.int64, device=xm.xla_device())
    tensor_in = torch.arange(
        world_size * world_size, dtype=torch.int64, device=xm.xla_device())
    dist.reduce_scatter(tensor_out, [tensor_in], op=dist.ReduceOp.SUM)
    xm.mark_step()

    torch.testing.assert_close(tensor_out.cpu(), expected)
//  [python]
def __triton_layernorm(self, dtype):
        shape = (20, 1024, 1024)
        origin_model = SimpleLN(shape[1:])
        apex_model = SimpleLN(shape[1:], layernorm_cls=ApexLayerNorm)
        apex_model.load_state_dict(origin_model.state_dict())
        atorch_model = SimpleLN(shape[1:], layernorm_cls=AtorchLayerNorm)
        atorch_model.load_state_dict(origin_model.state_dict())
        origin_model.to(dtype)
        apex_model.to(dtype)
        atorch_model.to(dtype)
        input_tensor = torch.randn(*shape, dtype=dtype, device="cuda")

        out_origin = origin_model(input_tensor)
        out_apex = apex_model(input_tensor)
        out_atorch = atorch_model(input_tensor)

        torch.testing.assert_allclose(out_origin, out_atorch)
        torch.testing.assert_allclose(out_origin, out_apex)
//  [python]
def __triton_layernorm(self, dtype):
        shape = (20, 1024, 1024)
        origin_model = SimpleLN(shape[1:])
        apex_model = SimpleLN(shape[1:], layernorm_cls=ApexLayerNorm)
        apex_model.load_state_dict(origin_model.state_dict())
        atorch_model = SimpleLN(shape[1:], layernorm_cls=AtorchLayerNorm)
        atorch_model.load_state_dict(origin_model.state_dict())
        origin_model.to(dtype)
        apex_model.to(dtype)
        atorch_model.to(dtype)
        input_tensor = torch.randn(*shape, dtype=dtype, device="cuda")

        out_origin = origin_model(input_tensor)
        out_apex = apex_model(input_tensor)
        out_atorch = atorch_model(input_tensor)

        torch.testing.assert_close(out_origin, out_atorch)
        torch.testing.assert_close(out_origin, out_apex)
//  [python]
def assert_all_close(
    actual: torch.Tensor,
    expected: torch.Tensor,
    base: torch.Tensor | None = None,
    rtol: float | None = None,
    atol: float | None = None,
    equal_nan: bool = False,
) -> None:
    if base is not None:
        actual = actual - base
        expected = expected - base

    if rtol is None or atol is None:
        from torch.testing._comparison import get_tolerances

        rtol, atol = get_tolerances(actual, expected, rtol=rtol, atol=atol)
        rtol *= 5 * NUM_UPDATES
        atol *= 5 * NUM_UPDATES

    t.assert_close(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        equal_nan=equal_nan,
        check_dtype=True,
    )
//  [python]
def assert_all_close(
    actual: torch.Tensor,
    expected: torch.Tensor,
    base: torch.Tensor | None = None,
    rtol: float | None = None,
    atol: float | None = None,
    equal_nan: bool = False,
) -> None:
    if base is not None:
        actual = actual - base
        expected = expected - base

    if rtol is None or atol is None:
        from torch.testing._comparison import get_tolerances

        rtol, atol = get_tolerances(actual, expected, rtol=rtol, atol=atol)
        rtol *= 5 * NUM_UPDATES
        atol *= 5 * NUM_UPDATES

    torch.testing.assert_close(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        equal_nan=equal_nan,
        check_dtype=True,
    )
//  [python]
def test_module_to_torchscript(
    m:torch.nn.Module, # The PyTorch module to be tested.
    inputs:Tensor, # A tensor or tuple of tensors representing the inputs to the model.
    trace:bool=True, #  If `True`, attempts to trace the model. Defaults to `True`.
    script:bool=True, # If `True`, attempts to script the model. Defaults to `True`.
    serialize:bool=True, # If `True`, saves and loads the traced/scripted module to ensure it can be serialized. Defaults to `True`.
    verbose:bool=True, # If `True`, prints detailed information about the tracing and scripting process. Defaults to `True`.
):
    "Tests if a PyTorch module can be correctly traced or scripted and serialized"
    
    m = m.eval()
    m_name = m.__class__.__name__

    # Ensure inputs are in a tuple or list format
    inp_is_tuple = isinstance(inputs, (tuple, list))

    # Get the model's output
    output = m(*inputs) if inp_is_tuple else m(inputs)
    output_shapes = output.shape if not isinstance(output, (tuple, list)) else [o.shape for o in output]
    if verbose:
        print(f"output.shape: {output_shapes}")

    # Try tracing the model
    if trace:
        if verbose:
            print("Tracing...")
        try:
            traced_m = torch.jit.trace(m, inputs)
            if serialize:
                file_path = Path(f"test_traced_{m_name}.pt")
                torch.jit.save(traced_m, file_path)
                traced_mod = torch.jit.load(file_path)
                file_path.unlink()
            traced_output = traced_m(*inputs) if inp_is_tuple else traced_m(inputs)
            torch.testing.assert_allclose(traced_output, output)
            if verbose:
                print(f"...{m_name} has been successfully traced 😃\n")
            return True
        except Exception as e:
            if verbose:
                print(f"{m_name} cannot be traced 😔")
                print(e)
                print("\n")

    # Try scripting the model
    if script:
        if verbose:
            print("Scripting...")
        try:
            scripted_m = torch.jit.script(m)
            if serialize:
                file_path = Path(f"test_scripted_{m_name}.pt")
                torch.jit.save(scripted_m, file_path)
                scripted_mod = torch.jit.load(file_path)
                file_path.unlink()
            scripted_output = scripted_m(*inputs) if inp_is_tuple else scripted_m(inputs)
            torch.testing.assert_allclose(scripted_output, output)
            if verbose:
                print(f"...{m_name} has been successfully scripted 😃\n")
            return True
        except Exception as e:
            if verbose:
                print(f"{m_name} cannot be scripted 😔")
                print(e)

    return False
//  [python]
def test_module_to_torchscript(
    m:torch.nn.Module, # The PyTorch module to be tested.
    inputs:Tensor, # A tensor or tuple of tensors representing the inputs to the model.
    trace:bool=True, #  If `True`, attempts to trace the model. Defaults to `True`.
    script:bool=True, # If `True`, attempts to script the model. Defaults to `True`.
    serialize:bool=True, # If `True`, saves and loads the traced/scripted module to ensure it can be serialized. Defaults to `True`.
    verbose:bool=True, # If `True`, prints detailed information about the tracing and scripting process. Defaults to `True`.
):
    "Tests if a PyTorch module can be correctly traced or scripted and serialized"
    
    m = m.eval()
    m_name = m.__class__.__name__

    # Ensure inputs are in a tuple or list format
    inp_is_tuple = isinstance(inputs, (tuple, list))

    # Get the model's output
    output = m(*inputs) if inp_is_tuple else m(inputs)
    output_shapes = output.shape if not isinstance(output, (tuple, list)) else [o.shape for o in output]
    if verbose:
        print(f"output.shape: {output_shapes}")

    # Try tracing the model
    if trace:
        if verbose:
            print("Tracing...")
        try:
            traced_m = torch.jit.trace(m, inputs)
            if serialize:
                file_path = Path(f"test_traced_{m_name}.pt")
                torch.jit.save(traced_m, file_path)
                traced_mod = torch.jit.load(file_path)
                file_path.unlink()
            traced_output = traced_m(*inputs) if inp_is_tuple else traced_m(inputs)
            torch.testing.assert_close(traced_output, output)
            if verbose:
                print(f"...{m_name} has been successfully traced 😃\n")
            return True
        except Exception as e:
            if verbose:
                print(f"{m_name} cannot be traced 😔")
                print(e)
                print("\n")

    # Try scripting the model
    if script:
        if verbose:
            print("Scripting...")
        try:
            scripted_m = torch.jit.script(m)
            if serialize:
                file_path = Path(f"test_scripted_{m_name}.pt")
                torch.jit.save(scripted_m, file_path)
                scripted_mod = torch.jit.load(file_path)
                file_path.unlink()
            scripted_output = scripted_m(*inputs) if inp_is_tuple else scripted_m(inputs)
            torch.testing.assert_close(scripted_output, output)
            if verbose:
                print(f"...{m_name} has been successfully scripted 😃\n")
            return True
        except Exception as e:
            if verbose:
                print(f"{m_name} cannot be scripted 😔")
                print(e)

    return False
//  [python]
group_norm_forward = create_group_norm_forward()
group_norm_silu_forward = create_group_norm_forward(act=activation.silu)

if __name__ == '__main__':
    import time
    import torch.nn.functional as F
    import torch as t

    def test_group_norm():
        x = torch.randn(2, 320, 32, 32).cuda().half()
        groups = 32
        weight = torch.randn(x.shape[1]).cuda().half()
        beta = torch.randn(x.shape[1]).cuda().half()

        torch.cuda.synchronize()
        begin_time = time.time()
        y = F.group_norm(x, groups, weight, beta, 1e-5)
        torch.cuda.synchronize()
        print('torch time: ', time.time() - begin_time)

        torch.cuda.synchronize()
        begin_time = time.time()
        y_triton = group_norm_forward(x, groups, weight, beta, 1e-5)[0]
        torch.cuda.synchronize()
        print('triton time: ', time.time() - begin_time)
        t.testing.assert_close(y_triton, y, rtol=1e-2, atol=1e-2)

        x_channel_last = x.contiguous(memory_format=torch.channels_last)
        torch.cuda.synchronize()
        begin_time = time.time()
        y_channel_last = F.group_norm(x_channel_last, groups, weight, beta,
                                      1e-5)
        y_channel_last = y_channel_last.contiguous(
            memory_format=torch.channels_last)
        torch.cuda.synchronize()
        print('torch channels last time: ', time.time() - begin_time)

        torch.cuda.synchronize()
        begin_time = time.time()
        y_triton_channel_last = group_norm_forward(x_channel_last, groups,
                                                   weight, beta, 1e-5)[0]
        torch.cuda.synchronize()
        print('triton channels last time: ', time.time() - begin_time)
        t.testing.assert_close(y_triton_channel_last,
                                   y_channel_last,
                                   rtol=1e-2,
                                   atol=1e-2)

    test_group_norm()
    test_group_norm()
    test_group_norm()
//  [python]
group_norm_forward = create_group_norm_forward()
group_norm_silu_forward = create_group_norm_forward(act=activation.silu)

if __name__ == '__main__':
    import time
    import torch.nn.functional as F

    def test_group_norm():
        x = torch.randn(2, 320, 32, 32).cuda().half()
        groups = 32
        weight = torch.randn(x.shape[1]).cuda().half()
        beta = torch.randn(x.shape[1]).cuda().half()

        torch.cuda.synchronize()
        begin_time = time.time()
        y = F.group_norm(x, groups, weight, beta, 1e-5)
        torch.cuda.synchronize()
        print('torch time: ', time.time() - begin_time)

        torch.cuda.synchronize()
        begin_time = time.time()
        y_triton = group_norm_forward(x, groups, weight, beta, 1e-5)[0]
        torch.cuda.synchronize()
        print('triton time: ', time.time() - begin_time)
        torch.testing.assert_close(y_triton, y, rtol=1e-2, atol=1e-2)

        x_channel_last = x.contiguous(memory_format=torch.channels_last)
        torch.cuda.synchronize()
        begin_time = time.time()
        y_channel_last = F.group_norm(x_channel_last, groups, weight, beta,
                                      1e-5)
        y_channel_last = y_channel_last.contiguous(
            memory_format=torch.channels_last)
        torch.cuda.synchronize()
        print('torch channels last time: ', time.time() - begin_time)

        torch.cuda.synchronize()
        begin_time = time.time()
        y_triton_channel_last = group_norm_forward(x_channel_last, groups,
                                                   weight, beta, 1e-5)[0]
        torch.cuda.synchronize()
        print('triton channels last time: ', time.time() - begin_time)
        torch.testing.assert_close(y_triton_channel_last,
                                   y_channel_last,
                                   rtol=1e-2,
                                   atol=1e-2)

    test_group_norm()
    test_group_norm()
    test_group_norm()
//  [python]
def test_clip_grad_norm(self):

        torch.manual_seed(self.seed + self.rank)

        # Identical models with data-parallel and ZeRO
        ref_model, ref_optim, dist_model, dist_optim = make_models(1, 1)

        # Training steps with pre-determined gradients
        xs = [3, 1, 4, 1, 5, 9]
        dys = [1, -1, 1, -1, 1, -1]
        for x, dy in zip(xs, dys):
            x = torch.tensor([[x]], dtype=torch.float32, device='cuda')
            dy = torch.tensor([[dy]], dtype=torch.float32, device='cuda')

            # Reference implementation
            ref_optim.zero_grad()
            y_ref = ref_model(x.detach())
            y_ref.backward(dy.detach())
            ref_grad_norm = torch.nn.utils.clip_grad_norm_(ref_model.parameters(), 3.5)
            ref_optim.step()

            # Distributed implementation
            dist_optim.zero_grad()
            y_dist = dist_model(x.detach())
            y_dist.backward(dy.detach())
            dist_grad_norm = dist_optim.clip_grad_norm(3.5)
            dist_optim.step()

            # Check that parameters match
            torch.testing.assert_allclose(dist_grad_norm, ref_grad_norm)
            for ref_param, dist_param in zip(ref_model.parameters(),
                                             dist_model.parameters()):
                torch.testing.assert_allclose(dist_param, ref_param)
//  [python]
def test_clip_grad_norm(self):

        torch.manual_seed(self.seed + self.rank)

        # Identical models with data-parallel and ZeRO
        ref_model, ref_optim, dist_model, dist_optim = make_models(1, 1)

        # Training steps with pre-determined gradients
        xs = [3, 1, 4, 1, 5, 9]
        dys = [1, -1, 1, -1, 1, -1]
        for x, dy in zip(xs, dys):
            x = torch.tensor([[x]], dtype=torch.float32, device='cuda')
            dy = torch.tensor([[dy]], dtype=torch.float32, device='cuda')

            # Reference implementation
            ref_optim.zero_grad()
            y_ref = ref_model(x.detach())
            y_ref.backward(dy.detach())
            ref_grad_norm = torch.nn.utils.clip_grad_norm_(ref_model.parameters(), 3.5)
            ref_optim.step()

            # Distributed implementation
            dist_optim.zero_grad()
            y_dist = dist_model(x.detach())
            y_dist.backward(dy.detach())
            dist_grad_norm = dist_optim.clip_grad_norm(3.5)
            dist_optim.step()

            # Check that parameters match
            torch.testing.assert_close(dist_grad_norm, ref_grad_norm)
            for ref_param, dist_param in zip(ref_model.parameters(),
                                             dist_model.parameters()):
                torch.testing.assert_close(dist_param, ref_param)
//  [python]
def _test_binary_precision_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        threshold: float = 0.5,
    ) -> None:

        input_np = np.where(input.numpy() < threshold, 0, 1)
        target_np = target.squeeze().numpy()
        sklearn_result = torch.tensor(precision_score(target_np, input_np)).to(
            torch.float32
        )

        torch.testing.assert_allclose(
            binary_precision(input, target, threshold=threshold),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_binary_precision_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        threshold: float = 0.5,
    ) -> None:

        input_np = np.where(input.numpy() < threshold, 0, 1)
        target_np = target.squeeze().numpy()
        sklearn_result = torch.tensor(precision_score(target_np, input_np)).to(
            torch.float32
        )

        torch.testing.assert_close(
            binary_precision(input, target, threshold=threshold),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_resnet_block(input_channels, output_channels, use_conv_shortcut, dropout_probability):
    block = ResNetBlock(input_channels=input_channels,
                        output_channels=output_channels,
                        use_conv_shortcut=use_conv_shortcut,
                        dropout_probability=dropout_probability,
                        zero_init_last=True)
    x = torch.randn(1, input_channels, 5, 5)
    y = block(x)
    assert y.shape == (1, output_channels, 5, 5), f'{y.shape}'
    if input_channels == output_channels:
        # Model should be exactly identity here
        torch.testing.assert_allclose(x, y)