jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_vectorized_error():
    rng = np.random.default_rng(0)
    generator = torch.Generator()
    generator.manual_seed(0)
    for _ in range(20):
        dim = rng.choice([1, 2])
        objective = th.Objective()
        batch_size = rng.choice(range(1, 11))

        n_vars = rng.choice([1, 10])
        vectors = [
            th.Vector(
                tensor=torch.randn(batch_size, dim, generator=generator), name=f"v{i}"
            )
            for i in range(n_vars)
        ]
        target = th.Vector(dim, name="target")
        w = th.ScaleCostWeight(torch.randn(1, generator=generator))
        for v in vectors:
            objective.add(th.Difference(v, target, w))

        se3s = [
            th.SE3(
                tensor=th.SE3.rand(batch_size, generator=generator).tensor,
                strict_checks=False,
            )
            for i in range(rng.choice([1, 10]))
        ]
        s_target = th.SE3.rand(1, generator=generator)
        ws = th.DiagonalCostWeight(torch.randn(1, 6, generator=generator))
        for s in se3s:
            objective.add(th.Difference(s, s_target, ws))

        vectorization = th.Vectorize(objective)
        objective.update_vectorization_if_needed()

        assert objective._vectorized_jacobians_iter is vectorization._cost_fn_wrappers
        _check_vectorized_wrappers(vectorization, objective)

        squared_error = torch.cat(
            [cf.weighted_error() for cf in objective.cost_functions.values()], dim=1
        )
        torch.testing.assert_close(squared_error, objective.error())

        # Check that calling error(also_update=False) changes the result but doesn't
        # change the vectorized wrappers
        another_squared_error = objective.error(
            {f"v{i}": torch.ones(batch_size, dim) for i in range(n_vars)},
            also_update=False,
        )
        assert not another_squared_error.allclose(squared_error)
        _check_vectorized_wrappers(vectorization, objective)

        # Just to be sure, check that objective.error() is calling
        # the vectorized error iterator.
        called = [False]

        def mock_err_iter(self):
            called[0] = True
            return self._cost_fn_wrappers  # just to have the same return type

        with mock.patch.object(
            th.Vectorize, "_get_vectorized_error_iter", mock_err_iter
        ):
            th.Vectorize(objective)
            objective.error()
            assert called[0]
//  [python]
def test_jagged_dense_bmm(
        self,
        B: int,
        M: int,
        N: int,
        max_L: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> None:
        assume(B != 0)
        torch.backends.cuda.matmul.allow_tf32 = False
        lengths = torch.randint(max_L + 1, size=(B,), device=device)
        total_length = int(lengths.sum().item())
        offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
        x_values = torch.rand(
            (total_length, M), requires_grad=True, dtype=dtype, device=device
        )
        y = torch.rand((B, M, N), requires_grad=True, dtype=dtype, device=device)

        output, _ = torch.ops.fbgemm.jagged_dense_bmm(
            x_values,
            offsets,
            y,
            max_L,
        )

        x_values_ref = x_values.detach().clone().requires_grad_(True)
        x_dense_ref = torch.ops.fbgemm.jagged_to_padded_dense(
            x_values_ref,
            [offsets],
            max_lengths=[max_L],
        )
        y_ref = y.detach().clone().requires_grad_(True)
        output_dense = torch.bmm(x_dense_ref, y_ref)

        output_ref, _ = torch.ops.fbgemm.dense_to_jagged(
            output_dense, [offsets], total_length
        )

        # verify forward
        torch.testing.assert_allclose(output, output_ref)
        # verify backward
        grad_output = output.detach().clone().requires_grad_(True)

        output.backward(grad_output)
        output_ref.backward(grad_output)

        torch.testing.assert_allclose(x_values.grad, x_values_ref.grad)
        torch.testing.assert_allclose(y.grad, y_ref.grad)
//  [python]
def test_jagged_dense_bmm(
        self,
        B: int,
        M: int,
        N: int,
        max_L: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> None:
        assume(B != 0)
        torch.backends.cuda.matmul.allow_tf32 = False
        lengths = torch.randint(max_L + 1, size=(B,), device=device)
        total_length = int(lengths.sum().item())
        offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
        x_values = torch.rand(
            (total_length, M), requires_grad=True, dtype=dtype, device=device
        )
        y = torch.rand((B, M, N), requires_grad=True, dtype=dtype, device=device)

        output, _ = torch.ops.fbgemm.jagged_dense_bmm(
            x_values,
            offsets,
            y,
            max_L,
        )

        x_values_ref = x_values.detach().clone().requires_grad_(True)
        x_dense_ref = torch.ops.fbgemm.jagged_to_padded_dense(
            x_values_ref,
            [offsets],
            max_lengths=[max_L],
        )
        y_ref = y.detach().clone().requires_grad_(True)
        output_dense = torch.bmm(x_dense_ref, y_ref)

        output_ref, _ = torch.ops.fbgemm.dense_to_jagged(
            output_dense, [offsets], total_length
        )

        # verify forward
        torch.testing.assert_close(output, output_ref)
        # verify backward
        grad_output = output.detach().clone().requires_grad_(True)

        output.backward(grad_output)
        output_ref.backward(grad_output)

        torch.testing.assert_close(x_values.grad, x_values_ref.grad)
        torch.testing.assert_close(y.grad, y_ref.grad)
//  [python]
def test_binary_accuracy_with_bool_target(self) -> None:
        num_classes = 2
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))

        output = binary_accuracy(input, target)
        output_bool = binary_accuracy(input, target.bool())

        torch.testing.assert_allclose(output, output_bool)
//  [python]
def test_binary_accuracy_with_bool_target(self) -> None:
        num_classes = 2
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))

        output = binary_accuracy(input, target)
        output_bool = binary_accuracy(input, target.bool())

        torch.testing.assert_close(output, output_bool)
//  [python]
def _test_binary_accuracy_with_input(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> None:
        input_np = input.numpy().round()
        target_np = target.squeeze().numpy()
        sklearn_result = torch.tensor(accuracy_score(target_np, input_np)).to(
            torch.float32
        )

        torch.testing.assert_allclose(
            binary_accuracy(input, target),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_binary_accuracy_with_input(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> None:
        input_np = input.numpy().round()
        target_np = target.squeeze().numpy()
        sklearn_result = torch.tensor(accuracy_score(target_np, input_np)).to(
            torch.float32
        )

        torch.testing.assert_close(
            binary_accuracy(input, target),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_rope(lit_llama, orig_llama) -> None:
    torch.manual_seed(1)

    bs, seq_len, n_head, n_embed = 1, 6, 2, 8
    x = torch.randint(0, 10000, size=(bs, seq_len, n_head, n_embed // n_head)).float()

    freqs_cis = orig_llama.precompute_freqs_cis(n_embed // n_head, seq_len)
    llama_rope_cache = lit_llama.build_rope_cache(seq_len, n_embed // n_head, dtype=x.dtype, device=x.device)
    torch.testing.assert_allclose(freqs_cis, torch.view_as_complex(llama_rope_cache))

    llama_x_rope = lit_llama.apply_rope(x, llama_rope_cache)
    orig_llama_x_rope, _ = orig_llama.apply_rotary_emb(x, x, freqs_cis)
    torch.testing.assert_allclose(llama_x_rope, orig_llama_x_rope)
//  [python]
def test_rope(lit_llama, orig_llama) -> None:
    torch.manual_seed(1)

    bs, seq_len, n_head, n_embed = 1, 6, 2, 8
    x = torch.randint(0, 10000, size=(bs, seq_len, n_head, n_embed // n_head)).float()

    freqs_cis = orig_llama.precompute_freqs_cis(n_embed // n_head, seq_len)
    llama_rope_cache = lit_llama.build_rope_cache(seq_len, n_embed // n_head, dtype=x.dtype, device=x.device)
    torch.testing.assert_close(freqs_cis, torch.view_as_complex(llama_rope_cache))

    llama_x_rope = lit_llama.apply_rope(x, llama_rope_cache)
    orig_llama_x_rope, _ = orig_llama.apply_rotary_emb(x, x, freqs_cis)
    torch.testing.assert_close(llama_x_rope, orig_llama_x_rope)
//  [python]
def _test_exact_match_with_input(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> None:
        input_np = input.numpy().round()
        target_np = target.squeeze().numpy()
        sklearn_result = torch.tensor(accuracy_score(target_np, input_np)).to(
            torch.float32
        )

        torch.testing.assert_allclose(
            multilabel_accuracy(input, target),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_exact_match_with_input(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> None:
        input_np = input.numpy().round()
        target_np = target.squeeze().numpy()
        sklearn_result = torch.tensor(accuracy_score(target_np, input_np)).to(
            torch.float32
        )

        torch.testing.assert_close(
            multilabel_accuracy(input, target),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_hamming_with_input(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> None:
        input_np = input.numpy().round()
        target_np = target.numpy()
        sklearn_result = torch.tensor(
            accuracy_score(target_np.flatten(), input_np.flatten())
        ).to(torch.float32)

        torch.testing.assert_allclose(
            multilabel_accuracy(input, target, criteria="hamming"),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_hamming_with_input(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> None:
        input_np = input.numpy().round()
        target_np = target.numpy()
        sklearn_result = torch.tensor(
            accuracy_score(target_np.flatten(), input_np.flatten())
        ).to(torch.float32)

        torch.testing.assert_close(
            multilabel_accuracy(input, target, criteria="hamming"),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_accuracy_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
    ) -> None:
        target_np = target.squeeze().flatten().numpy()
        input_squeezed = input.squeeze()
        input_label_ids = (
            torch.argmax(input_squeezed, dim=1)
            if input_squeezed.ndim == 2
            else input_squeezed
        )
        input_np = input_label_ids.flatten().numpy()
        compute_result = torch.tensor(accuracy_score(target_np, input_np)).to(
            torch.float32
        )
        torch.testing.assert_allclose(
            multiclass_accuracy(input, target),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_accuracy_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
    ) -> None:
        target_np = target.squeeze().flatten().numpy()
        input_squeezed = input.squeeze()
        input_label_ids = (
            torch.argmax(input_squeezed, dim=1)
            if input_squeezed.ndim == 2
            else input_squeezed
        )
        input_np = input_label_ids.flatten().numpy()
        compute_result = torch.tensor(accuracy_score(target_np, input_np)).to(
            torch.float32
        )
        torch.testing.assert_close(
            multiclass_accuracy(input, target),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_accuracy_class_average(self) -> None:
        num_classes = 4
        # high=num_classes-1 gives us NaN value for the last class
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes - 1, size=(BATCH_SIZE,))

        input_flattened = input.flatten()
        target_flattened = target.flatten()
        accuracy_per_class = np.empty(num_classes)
        for i in range(num_classes):
            accuracy_per_class[i] = accuracy_score(
                target_flattened[target_flattened == i].numpy(),
                input_flattened[target_flattened == i].numpy(),
            )

        self.assertEqual(
            multiclass_accuracy(
                input, target, average="macro", num_classes=num_classes
            ),
            torch.tensor(np.mean(accuracy_per_class[~np.isnan(accuracy_per_class)])).to(
                torch.float32
            ),
        )

        self.assertEqual(
            multiclass_accuracy(input, target, average=None, num_classes=num_classes),
            torch.tensor(accuracy_per_class).view(-1).to(torch.float32),
            equal_nan=True,
        )
//  [python]
def test_accuracy_class_average(self) -> None:
        num_classes = 4
        # high=num_classes-1 gives us NaN value for the last class
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes - 1, size=(BATCH_SIZE,))

        input_flattened = input.flatten()
        target_flattened = target.flatten()
        accuracy_per_class = np.empty(num_classes)
        for i in range(num_classes):
            accuracy_per_class[i] = accuracy_score(
                target_flattened[target_flattened == i].numpy(),
                input_flattened[target_flattened == i].numpy(),
            )

        torch.testing.assert_close(
            multiclass_accuracy(
                input, target, average="macro", num_classes=num_classes
            ),
            torch.tensor(np.mean(accuracy_per_class[~np.isnan(accuracy_per_class)])).to(
                torch.float32
            ),
            atol=1e-8,
            rtol=1e-5,
        )

        torch.testing.assert_close(
            multiclass_accuracy(input, target, average=None, num_classes=num_classes),
            torch.tensor(accuracy_per_class).view(-1).to(torch.float32),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_reward_scaling(self, batch, scale, loc, keys, device, standard_normal):
        torch.manual_seed(0)
        if keys is None:
            keys_total = set()
        else:
            keys_total = set(keys)
        reward_scaling = RewardScaling(
            in_keys=keys, scale=scale, loc=loc, standard_normal=standard_normal
        )
        td = TensorDict(
            {
                **{key: torch.randn(*batch, 1, device=device) for key in keys_total},
                "reward": torch.randn(*batch, 1, device=device),
            },
            batch,
            device=device,
        )
        td.set("dont touch", torch.randn(*batch, 1, device=device))
        td_copy = td.clone()
        reward_scaling(td)
        for key in keys_total:
            if standard_normal:
                original_key = td.get(key)
                scaled_key = (td_copy.get(key) - loc) / scale
                torch.testing.assert_allclose(original_key, scaled_key)
            else:
                original_key = td.get(key)
                scaled_key = td_copy.get(key) * scale + loc
                torch.testing.assert_allclose(original_key, scaled_key)
        if keys is None:
            if standard_normal:
                original_key = td.get("reward")
                scaled_key = (td_copy.get("reward") - loc) / scale
                torch.testing.assert_allclose(original_key, scaled_key)
            else:
                original_key = td.get("reward")
                scaled_key = td_copy.get("reward") * scale + loc
                torch.testing.assert_allclose(original_key, scaled_key)

        assert (td.get("dont touch") == td_copy.get("dont touch")).all()

        if len(keys_total) == 1:
            reward_spec = UnboundedContinuousTensorSpec(device=device)
            reward_spec = reward_scaling.transform_reward_spec(reward_spec)
            assert reward_spec.shape == torch.Size([1])
//  [python]
def test_reward_scaling(self, batch, scale, loc, keys, device, standard_normal):
        torch.manual_seed(0)
        if keys is None:
            keys_total = set()
        else:
            keys_total = set(keys)
        reward_scaling = RewardScaling(
            in_keys=keys, scale=scale, loc=loc, standard_normal=standard_normal
        )
        td = TensorDict(
            {
                **{key: torch.randn(*batch, 1, device=device) for key in keys_total},
                "reward": torch.randn(*batch, 1, device=device),
            },
            batch,
            device=device,
        )
        td.set("dont touch", torch.randn(*batch, 1, device=device))
        td_copy = td.clone()
        reward_scaling(td)
        for key in keys_total:
            if standard_normal:
                original_key = td.get(key)
                scaled_key = (td_copy.get(key) - loc) / scale
                torch.testing.assert_close(original_key, scaled_key)
            else:
                original_key = td.get(key)
                scaled_key = td_copy.get(key) * scale + loc
                torch.testing.assert_close(original_key, scaled_key)
        if keys is None:
            if standard_normal:
                original_key = td.get("reward")
                scaled_key = (td_copy.get("reward") - loc) / scale
                torch.testing.assert_close(original_key, scaled_key)
            else:
                original_key = td.get("reward")
                scaled_key = td_copy.get("reward") * scale + loc
                torch.testing.assert_close(original_key, scaled_key)

        assert (td.get("dont touch") == td_copy.get("dont touch")).all()

        if len(keys_total) == 1:
            reward_spec = UnboundedContinuousTensorSpec(device=device)
            reward_spec = reward_scaling.transform_reward_spec(reward_spec)
            assert reward_spec.shape == torch.Size([1])
//  [python]
def test_topk_accuracy_with_input(self) -> None:
        input = torch.tensor(
            [
                [0.9, 0.1, 0, 0],
                [0.1, 0.2, 0.4, 0.3],
                [0.4, 0.4, 0.1, 0.1],
                [0, 0, 0.8, 0.2],
            ]
        )
        target = torch.tensor([0, 1, 2, 3])

        compute_result = torch.tensor(0.5)
        t.assert_allclose(
            multiclass_accuracy(input, target, k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        compute_result = torch.tensor(1.0)
        t.assert_allclose(
            multiclass_accuracy(input, target, k=4),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_with_input(self) -> None:
        input = torch.tensor(
            [
                [0.9, 0.1, 0, 0],
                [0.1, 0.2, 0.4, 0.3],
                [0.4, 0.4, 0.1, 0.1],
                [0, 0, 0.8, 0.2],
            ]
        )
        target = torch.tensor([0, 1, 2, 3])

        compute_result = torch.tensor(0.5)
        torch.testing.assert_close(
            multiclass_accuracy(input, target, k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        compute_result = torch.tensor(1.0)
        torch.testing.assert_close(
            multiclass_accuracy(input, target, k=4),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_base(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(0.0)
        t.assert_allclose(
            topk_multilabel_accuracy(input, target, k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_base(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(0.0)
        torch.testing.assert_close(
            topk_multilabel_accuracy(input, target, k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_hamming(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(7 / 12)
        torch.testing.assert_allclose(
            topk_multilabel_accuracy(input, target, criteria="hamming", k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_hamming(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(7 / 12)
        torch.testing.assert_close(
            topk_multilabel_accuracy(input, target, criteria="hamming", k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_with_input_macro(self) -> None:
        input = torch.tensor(
            [
                [0.9, 0.1, 0, 0],
                [0.9, 0.1, 0, 0],
                [0.3, 0.2, 0.4, 0.1],
                [0.3, 0.2, 0.4, 0.1],
                [0.4, 0.4, 0.1, 0.1],
                [0.3, 0.5, 0.1, 0.1],
            ]
        )
        target = torch.tensor([0, 0, 0, 0, 2, 2])

        compute_result = torch.tensor(0.5)
        torch.testing.assert_allclose(
            multiclass_accuracy(input, target, average="macro", num_classes=4, k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        compute_result = torch.tensor(1.0)
        torch.testing.assert_allclose(
            multiclass_accuracy(input, target, average="macro", num_classes=4, k=4),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_with_input_macro(self) -> None:
        input = torch.tensor(
            [
                [0.9, 0.1, 0, 0],
                [0.9, 0.1, 0, 0],
                [0.3, 0.2, 0.4, 0.1],
                [0.3, 0.2, 0.4, 0.1],
                [0.4, 0.4, 0.1, 0.1],
                [0.3, 0.5, 0.1, 0.1],
            ]
        )
        target = torch.tensor([0, 0, 0, 0, 2, 2])

        compute_result = torch.tensor(0.5)
        torch.testing.assert_close(
            multiclass_accuracy(input, target, average="macro", num_classes=4, k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        compute_result = torch.tensor(1.0)
        torch.testing.assert_close(
            multiclass_accuracy(input, target, average="macro", num_classes=4, k=4),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_overlap(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(4 / 4)
        t.assert_close(
            topk_multilabel_accuracy(input, target, criteria="overlap", k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_overlap(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(4 / 4)
        torch.testing.assert_close(
            topk_multilabel_accuracy(input, target, criteria="overlap", k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_with_input_no_average(self) -> None:
        input = torch.tensor(
            [
                [0.9, 0.1, 0, 0],
                [0.9, 0.1, 0, 0],
                [0.3, 0.2, 0.4, 0.1],
                [0.3, 0.2, 0.4, 0.1],
                [0.4, 0.4, 0.1, 0.1],
                [0.3, 0.5, 0.1, 0.1],
            ]
        )
        target = torch.tensor([0, 0, 0, 0, 2, 2])

        compute_result = torch.tensor([1.0, np.NAN, 0, np.NAN])
        torch.testing.assert_allclose(
            multiclass_accuracy(input, target, average=None, num_classes=4, k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        compute_result = torch.tensor([1.0, np.NAN, 1.0, np.NAN])
        torch.testing.assert_allclose(
            multiclass_accuracy(input, target, average=None, num_classes=4, k=4),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_with_input_no_average(self) -> None:
        input = torch.tensor(
            [
                [0.9, 0.1, 0, 0],
                [0.9, 0.1, 0, 0],
                [0.3, 0.2, 0.4, 0.1],
                [0.3, 0.2, 0.4, 0.1],
                [0.4, 0.4, 0.1, 0.1],
                [0.3, 0.5, 0.1, 0.1],
            ]
        )
        target = torch.tensor([0, 0, 0, 0, 2, 2])

        compute_result = torch.tensor([1.0, np.NAN, 0, np.NAN])
        torch.testing.assert_close(
            multiclass_accuracy(input, target, average=None, num_classes=4, k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        compute_result = torch.tensor([1.0, np.NAN, 1.0, np.NAN])
        torch.testing.assert_close(
            multiclass_accuracy(input, target, average=None, num_classes=4, k=4),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_multilabel_accuracy_criteria(self) -> None:
        input = torch.tensor([[0, 1, 1], [1, 1, 0], [1, 1, 1], [0, 1, 1]])
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])
        # test overlap criteria with positive inputs
        overlap_accuracy = multilabel_accuracy(input, target, criteria="overlap")
        torch.testing.assert_allclose(
            overlap_accuracy,
            torch.tensor(1.0),
            atol=1e-8,
            rtol=1e-5,
        )
        # test overlap criteria with no positive inputs: input[2]
        input_overlap_no_pos = torch.tensor(
            [[0, 0, 1], [0, 0, 0], [0, 0, 0], [1, 0, 0]]
        )
        target_overlap_no_pos = torch.tensor(
            [[0, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 1]]
        )
        overlap_accuracy = multilabel_accuracy(
            input_overlap_no_pos, target_overlap_no_pos, criteria="overlap"
        )
        torch.testing.assert_allclose(
            overlap_accuracy,
            torch.tensor(1 / 4),
            atol=1e-8,
            rtol=1e-5,
        )
        # test contain criteria: input[1], input[2], input[3]
        contain_accuracy = multilabel_accuracy(input, target, criteria="contain")
        torch.testing.assert_allclose(
            contain_accuracy,
            torch.tensor(3 / 4),
            atol=1e-8,
            rtol=1e-5,
        )
        # test belong criteria: input[2]
        belong_accuracy = multilabel_accuracy(input, target, criteria="belong")
        torch.testing.assert_allclose(
            belong_accuracy,
            torch.tensor(1 / 4),
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_multilabel_accuracy_criteria(self) -> None:
        input = torch.tensor([[0, 1, 1], [1, 1, 0], [1, 1, 1], [0, 1, 1]])
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])
        # test overlap criteria with positive inputs
        overlap_accuracy = multilabel_accuracy(input, target, criteria="overlap")
        torch.testing.assert_close(
            overlap_accuracy,
            torch.tensor(1.0),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
        # test overlap criteria with no positive inputs: input[2]
        input_overlap_no_pos = torch.tensor(
            [[0, 0, 1], [0, 0, 0], [0, 0, 0], [1, 0, 0]]
        )
        target_overlap_no_pos = torch.tensor(
            [[0, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 1]]
        )
        overlap_accuracy = multilabel_accuracy(
            input_overlap_no_pos, target_overlap_no_pos, criteria="overlap"
        )
        torch.testing.assert_close(
            overlap_accuracy,
            torch.tensor(1 / 4),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
        # test contain criteria: input[1], input[2], input[3]
        contain_accuracy = multilabel_accuracy(input, target, criteria="contain")
        torch.testing.assert_close(
            contain_accuracy,
            torch.tensor(3 / 4),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
        # test belong criteria: input[2]
        belong_accuracy = multilabel_accuracy(input, target, criteria="belong")
        torch.testing.assert_close(
            belong_accuracy,
            torch.tensor(1 / 4),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_contain(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(2 / 4)
        torch.testing.assert_allclose(
            topk_multilabel_accuracy(input, target, criteria="contain", k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_contain(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(2 / 4)
        torch.testing.assert_close(
            topk_multilabel_accuracy(input, target, criteria="contain", k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_belong(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(1 / 4)
        torch.testing.assert_allclose(
            topk_multilabel_accuracy(input, target, criteria="belong", k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_topk_accuracy_belong(self) -> None:
        input = torch.tensor(
            [[0.1, 0.5, 0.2], [0.3, 0.2, 0.1], [0.2, 0.4, 0.5], [0, 0.1, 0.9]]
        )
        target = torch.tensor([[1, 1, 0], [0, 1, 0], [1, 1, 1], [0, 1, 0]])

        compute_result = torch.tensor(1 / 4)
        torch.testing.assert_close(
            topk_multilabel_accuracy(input, target, criteria="belong", k=2),
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_vee(batch_size: Union[int, Sequence[int]], dtype: torch.dtype):
    if isinstance(batch_size, int):
        batch_size = (batch_size,)

    rng = torch.Generator()
    rng.manual_seed(0)
    tangent_vector = torch.rand(*batch_size, 6, dtype=dtype, generator=rng)
    matrix = se3_impl._hat_autograd_fn(tangent_vector)

    # check analytic backward for the operator
    check_lie_group_function(se3_impl, "vee", TEST_EPS, (matrix,))

    # check the correctness of hat and vee
    actual_tangent_vector = se3_impl._vee_autograd_fn(matrix)
    t.assert_close(
        actual_tangent_vector, tangent_vector, atol=TEST_EPS, rtol=TEST_EPS
    )
//  [python]
def test_vee(batch_size: Union[int, Sequence[int]], dtype: torch.dtype):
    if isinstance(batch_size, int):
        batch_size = (batch_size,)

    rng = torch.Generator()
    rng.manual_seed(0)
    tangent_vector = torch.rand(*batch_size, 6, dtype=dtype, generator=rng)
    matrix = se3_impl._hat_autograd_fn(tangent_vector)

    # check analytic backward for the operator
    check_lie_group_function(se3_impl, "vee", TEST_EPS, (matrix,))

    # check the correctness of hat and vee
    actual_tangent_vector = se3_impl._vee_autograd_fn(matrix)
    torch.testing.assert_close(
        actual_tangent_vector, tangent_vector, atol=TEST_EPS, rtol=TEST_EPS
    )
//  [python]
def test_vanilla_replacement(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv1(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, 8, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_vanilla_replacement(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv1(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 8, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_module_list_replacement(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.ModuleList([nn.Conv2d(3, 8, kernel_size=3, padding=1) for _ in range(3)])

            def forward(self, x):
                return torch.cat([conv(x) for conv in self.convs], dim=1)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, 3 * 8, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_module_list_replacement(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.ModuleList([nn.Conv2d(3, 8, kernel_size=3, padding=1) for _ in range(3)])

            def forward(self, x):
                return torch.cat([conv(x) for conv in self.convs], dim=1)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 3 * 8, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_sequential_list_replacement(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.Sequential(
                    nn.Conv2d(3, 8, kernel_size=3, padding=1),
                    nn.Conv2d(8, 16, kernel_size=3, padding=1),
                )

            def forward(self, x):
                return self.convs(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, 16, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_sequential_list_replacement(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.Sequential(
                    nn.Conv2d(3, 8, kernel_size=3, padding=1),
                    nn.Conv2d(8, 16, kernel_size=3, padding=1),
                )

            def forward(self, x):
                return self.convs(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 16, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_static_selective_skip_quantization(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer, SkipQuantization

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = SkipQuantization(nn.Conv2d(8, 16, kernel_size=3, padding=1))

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, 16, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.conv2, nn.Conv2d))
//  [python]
def test_static_selective_skip_quantization(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer, SkipQuantization

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = SkipQuantization(nn.Conv2d(8, 16, kernel_size=3, padding=1))

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 16, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.conv2, nn.Conv2d))
//  [python]
def test_nested_module_replacement(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block.linear, SelectiveQuantizer.mapping_instructions[nn.Linear].quantized_target_class))
//  [python]
def test_nested_module_replacement(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block.linear, SelectiveQuantizer.mapping_instructions[nn.Linear].quantized_target_class))
//  [python]
def test_dynamic_skip_quantization(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.register_skip_quantization(layer_names={"conv2"})
        q_util.quantize_module(module)
        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertTrue(torch.equal(y.size(), torch.Size([1, 16, 32, 32])))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.conv2, nn.Conv2d))
//  [python]
def test_dynamic_skip_quantization(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.register_skip_quantization(layer_names={"conv2"})
        q_util.quantize_module(module)
        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 16, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.conv2, nn.Conv2d))
//  [python]
def test_non_default_quant_descriptors_are_piped(self):
        from super_gradients.training.utils.quantization import SelectiveQuantizer
        from pytorch_quantization.calib import MaxCalibrator

        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv1(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer(default_quant_modules_calibrator_inputs="max", default_quant_modules_calibrator_weights="max")
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_allclose(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == MaxCalibrator)