def test_torch_ref_match():

    # Verify if the torch implementation values match the original Numpy implementation.

    num_teachers, num_examples, num_labels = (100, 50, 10)
    preds = (np.random.rand(num_teachers, num_examples) * num_labels).astype(int)  # fake preds

    indices = (np.random.rand(num_examples) * num_labels).astype(int)  # true answers

    preds[:, 0:10] *= 0

    data_dep_eps, data_ind_eps = pate.perform_analysis_torch(
        preds, indices, noise_eps=0.1, delta=1e-5
    )

    data_dep_eps_ref, data_ind_eps_ref = pate.perform_analysis(
        preds, indices, noise_eps=0.1, delta=1e-5
    )

    assert torch.isclose(data_dep_eps, torch.tensor(data_dep_eps_ref, dtype=torch.float32))
    assert torch.isclose(data_ind_eps, torch.tensor(data_ind_eps_ref, dtype=torch.float32)) 

def test_hinge_loss():
    true_value = torch.Tensor([[1.2], [1], [1], [1]])
    pred_value = torch.Tensor([[1.2], [0.1], [0], [-0.3]])
    expected_loss = torch.Tensor([(0 + 1 - 0.3 + 0) / 2.0])
    loss = losses.RankHingeLoss()(pred_value, true_value)
    assert torch.isclose(expected_loss, loss)
    expected_loss = torch.Tensor(
        [(2 + 0.1 - 1.2 + 2 - 0.3 + 0) / 2.0])
    loss = losses.RankHingeLoss(margin=2)(pred_value, true_value)
    assert torch.isclose(expected_loss, loss)
    true_value = torch.Tensor(
        [[1.2], [1], [0.8], [1], [1], [0.8]])
    pred_value = torch.Tensor(
        [[1.2], [0.1], [-0.5], [0], [0], [-0.3]])
    expected_loss = torch.Tensor(
        [(0 + 1 - 0.15) / 2.0])
    loss = losses.RankHingeLoss(num_neg=2, margin=1)(
        pred_value, true_value)
    assert torch.isclose(expected_loss, loss) 

def test_rank_crossentropy_loss():
    losses.neg_num = 1

    def softmax(x):
        return np.exp(x) / np.sum(np.exp(x), axis=0)

    true_value = torch.Tensor([[1.], [0.], [0.], [1.]])
    pred_value = torch.Tensor([[0.8], [0.1], [0.8], [0.1]])
    expected_loss = torch.Tensor(
        [(-np.log(softmax([0.8, 0.1])[0]) - np.log(
            softmax([0.8, 0.1])[1])) / 2])
    loss = losses.RankCrossEntropyLoss()(pred_value, true_value)
    assert torch.isclose(expected_loss, loss)
    true_value = torch.Tensor([[1.], [0.], [0.], [0.], [1.], [0.]])
    pred_value = torch.Tensor([[0.8], [0.1], [0.1], [0.8], [0.1], [0.1]])
    expected_loss = torch.Tensor(
        [(-np.log(softmax([0.8, 0.1, 0.1])[0]) - np.log(
            softmax([0.8, 0.1, 0.1])[1])) / 2])
    loss = losses.RankCrossEntropyLoss(num_neg=2)(
        pred_value, true_value)
    assert torch.isclose(expected_loss, loss) 

def _components(self) -> Dict[Tuple[str, str, str], Tuple[Tensor, Tensor]]:
        states_per_measure = defaultdict(list)
        for state_belief in self.state_beliefs:
            for m, measure in enumerate(self.design.measures):
                H = state_belief.H[:, m, :].data
                m = H * state_belief.means.data
                std = H * torch.diagonal(state_belief.covs.data, dim1=-2, dim2=-1).sqrt()
                states_per_measure[measure].append((m, std))

        out = {}
        for measure, means_and_stds in states_per_measure.items():
            means, stds = zip(*means_and_stds)
            means = torch.stack(means).permute(1, 0, 2)
            stds = torch.stack(stds).permute(1, 0, 2)
            for s, (process_name, state_element) in enumerate(self.design.state_elements):
                if ~torch.isclose(means[:, :, s].abs().max(), torch.zeros(1)):
                    out[(measure, process_name, state_element)] = (means[:, :, s], stds[:, :, s])
        return out 

def test_rand(self):
        """Tests uniform random variable generation on [0, 1)"""
        for size in [(10,), (10, 10), (10, 10, 10)]:
            randvec = crypten.rand(*size)
            self.assertTrue(randvec.size() == size, "Incorrect size")
            tensor = randvec.get_plain_text()
            self.assertTrue(
                (tensor >= 0).all() and (tensor < 1).all(), "Invalid values"
            )

        randvec = crypten.rand(int(1e6)).get_plain_text()
        mean = torch.mean(randvec)
        var = torch.var(randvec)
        self.assertTrue(torch.isclose(mean, torch.Tensor([0.5]), rtol=1e-3, atol=1e-3))
        self.assertTrue(
            torch.isclose(var, torch.Tensor([1.0 / 12]), rtol=1e-3, atol=1e-3)
        ) 

def _run_act_layer_grad(act_type):
    x = torch.rand(10, 1000) * 10
    m = MLP(act_layer=act_type)

    def _run(x, act_layer=''):
        if act_layer:
            # replace act layer if set
            m.act = create_act_layer(act_layer, inplace=True)
        out = m(x)
        l = (out - 0).pow(2).sum()
        return l

    out_me = _run(x)

    with set_layer_config(scriptable=True):
        out_jit = _run(x, act_type)

    assert torch.isclose(out_jit, out_me)

    with set_layer_config(no_jit=True):
        out_basic = _run(x, act_type)

    assert torch.isclose(out_basic, out_jit) 

def test_arcface_loss(self):
        margin = 30
        scale = 64
        loss_func = ArcFaceLoss(margin=margin, scale=scale, num_classes=10, embedding_size=2)

        embedding_angles = torch.arange(0, 180)
        embeddings = torch.tensor([c_f.angle_to_coord(a) for a in embedding_angles], requires_grad=True, dtype=torch.float) #2D embeddings
        labels = torch.randint(low=0, high=10, size=(180,))

        loss = loss_func(embeddings, labels)
        loss.backward()

        weights = torch.nn.functional.normalize(loss_func.W, p=2, dim=0)
        logits = torch.matmul(embeddings, weights)
        for i, c in enumerate(labels):
            logits[i, c] = torch.cos(torch.acos(logits[i, c]) + torch.tensor(np.radians(margin)))
        
        correct_loss = torch.nn.functional.cross_entropy(logits*scale, labels)
        self.assertTrue(torch.isclose(loss, correct_loss)) 

def test_triplet_margin_loss(self):
        margin = 0.2
        loss_funcA = TripletMarginLoss(margin=margin)
        loss_funcB = TripletMarginLoss(margin=margin, reducer=MeanReducer())
        embedding_angles = [0, 20, 40, 60, 80]
        embeddings = torch.tensor([c_f.angle_to_coord(a) for a in embedding_angles], requires_grad=True, dtype=torch.float) #2D embeddings
        labels = torch.LongTensor([0, 0, 1, 1, 2])

        lossA = loss_funcA(embeddings, labels)
        lossB = loss_funcB(embeddings, labels)
        lossA.backward()
        lossB.backward()
        
        triplets = [(0,1,2), (0,1,3), (0,1,4), (1,0,2), (1,0,3), (1,0,4), (2,3,0), (2,3,1), (2,3,4), (3,2,0), (3,2,1), (3,2,4)]

        correct_loss = 0
        num_non_zero_triplets = 0
        for a, p, n in triplets:
            anchor, positive, negative = embeddings[a], embeddings[p], embeddings[n]
            curr_loss = torch.relu(torch.sqrt(torch.sum((anchor-positive)**2)) - torch.sqrt(torch.sum((anchor-negative)**2)) + margin)
            if curr_loss > 0:
                num_non_zero_triplets += 1
            correct_loss += curr_loss
        self.assertTrue(torch.isclose(lossA, correct_loss/num_non_zero_triplets))
        self.assertTrue(torch.isclose(lossB, correct_loss/len(triplets))) 

def test_angular_loss(self):
        loss_func = AngularLoss(alpha=40)
        embedding_angles = [0, 20, 40, 60, 80]
        embeddings = torch.tensor([c_f.angle_to_coord(a) for a in embedding_angles], requires_grad=True, dtype=torch.float) #2D embeddings
        labels = torch.LongTensor([0, 0, 1, 1, 2])

        loss = loss_func(embeddings, labels)
        loss.backward()
        sq_tan_alpha = torch.tan(torch.tensor(np.radians(40)))**2
        triplets = [(0,1,2), (0,1,3), (0,1,4), (1,0,2), (1,0,3), (1,0,4), (2,3,0), (2,3,1), (2,3,4), (3,2,0), (3,2,1), (3,2,4)]

        correct_losses = [0,0,0,0]
        for a, p, n in triplets:
            anchor, positive, negative = embeddings[a], embeddings[p], embeddings[n]
            exponent = 4*sq_tan_alpha*torch.matmul(anchor+positive,negative) - 2*(1+sq_tan_alpha)*torch.matmul(anchor, positive)
            correct_losses[a] += torch.exp(exponent)
        total_loss = 0
        for c in correct_losses:
            total_loss += torch.log(1+c)
        total_loss /= len(correct_losses)
        self.assertTrue(torch.isclose(loss, total_loss.to(torch.float32))) 

def test_cosface_loss(self):
        margin = 0.5
        scale = 64
        loss_func = CosFaceLoss(margin=margin, scale=scale, num_classes=10, embedding_size=2)

        embedding_angles = torch.arange(0, 180)
        embeddings = torch.tensor([c_f.angle_to_coord(a) for a in embedding_angles], requires_grad=True, dtype=torch.float) #2D embeddings
        labels = torch.randint(low=0, high=10, size=(180,))

        loss = loss_func(embeddings, labels)
        loss.backward()

        weights = torch.nn.functional.normalize(loss_func.W, p=2, dim=0)
        logits = torch.matmul(embeddings, weights)
        for i, c in enumerate(labels):
            logits[i, c] -= margin
        
        correct_loss = torch.nn.functional.cross_entropy(logits*scale, labels)
        self.assertTrue(torch.isclose(loss, correct_loss)) 

def test_proxy_nca_loss(self):
        softmax_scale = 10
        loss_func = ProxyNCALoss(softmax_scale=softmax_scale, num_classes=10, embedding_size=2)

        embedding_angles = torch.arange(0, 180)
        embeddings = torch.tensor([c_f.angle_to_coord(a) for a in embedding_angles], requires_grad=True, dtype=torch.float) #2D embeddings
        labels = torch.randint(low=0, high=10, size=(180,))

        loss = loss_func(embeddings, labels)
        loss.backward()

        proxies = torch.nn.functional.normalize(loss_func.proxies, p=2, dim=1)
        correct_loss = 0
        for i in range(len(embeddings)):
            curr_emb, curr_label = embeddings[i], labels[i]
            curr_proxy = proxies[curr_label]
            denominator = torch.sum((curr_emb-proxies)**2, dim=1)
            denominator = torch.sum(torch.exp(-denominator*softmax_scale))
            numerator = torch.sum((curr_emb-curr_proxy)**2)
            numerator = torch.exp(-numerator*softmax_scale)
            correct_loss += -torch.log(numerator/denominator)
        
        correct_loss /= len(embeddings)
        self.assertTrue(torch.isclose(loss, correct_loss)) 

def test_add_infinity_and_beyond(a, b, c, negative, manifold, dtype):
    _a = a
    if torch.isclose(c, c.new_zeros(())).any():
        pytest.skip("zero not checked")
    infty = b * 10000000
    for i in range(100):
        z = manifold.expmap(a, infty, project=False)
        z = manifold.projx(z)
        assert not torch.isnan(z).any(), ("Found nans", i, z)
        assert torch.isfinite(z).all(), ("Found Infs", i, z)
        z = manifold.mobius_scalar_mul(
            torch.tensor(1000.0, dtype=z.dtype), z, project=False
        )
        z = manifold.projx(z)
        assert not torch.isnan(z).any(), ("Found nans", i, z)
        assert torch.isfinite(z).all(), ("Found Infs", i, z)

        infty = manifold.transp(a, z, infty)
        assert torch.isfinite(infty).all(), (i, infty)
        a = z
    z = manifold.expmap(a, -infty)
    # they just need to be very far, exact answer is not supposed
    tolerance = {
        torch.float32: dict(rtol=3e-1, atol=2e-1),
        torch.float64: dict(rtol=1e-1, atol=1e-3),
    }
    if negative:
        np.testing.assert_allclose(z.detach(), -a.detach(), **tolerance[dtype])
    else:
        assert not torch.isnan(z).any(), "Found nans"
        assert not torch.isnan(a).any(), "Found nans" 

def test_weighted_midpoint(_k, lincomb):
    manifold = stereographic.Stereographic(_k, learnable=True)
    a = manifold.random(2, 3, 10).requires_grad_(True)
    mid = manifold.weighted_midpoint(a, lincomb=lincomb)
    assert torch.isfinite(mid).all()
    assert mid.shape == (a.shape[-1],)
    mid.sum().backward()
    assert torch.isfinite(a.grad).all()
    assert not torch.isclose(manifold.k.grad, manifold.k.new_zeros(())) 

def test_weighted_midpoint_reduce_dim(_k, lincomb):
    manifold = stereographic.Stereographic(_k, learnable=True)
    a = manifold.random(2, 3, 10).requires_grad_(True)
    mid = manifold.weighted_midpoint(a, reducedim=[0], lincomb=lincomb)
    assert mid.shape == a.shape[-2:]
    assert torch.isfinite(mid).all()
    mid.sum().backward()
    assert torch.isfinite(a.grad).all()
    assert not torch.isclose(manifold.k.grad, manifold.k.new_zeros(())) 

def test_weighted_midpoint_weighted(_k, lincomb):
    manifold = stereographic.Stereographic(_k, learnable=True)
    a = manifold.random(2, 3, 10).requires_grad_(True)
    mid = manifold.weighted_midpoint(
        a, reducedim=[0], lincomb=lincomb, weights=torch.rand_like(a[..., 0])
    )
    assert mid.shape == a.shape[-2:]
    assert torch.isfinite(mid).all()
    mid.sum().backward()
    assert torch.isfinite(a.grad).all()
    assert not torch.isclose(manifold.k.grad, manifold.k.new_zeros(())) 

def test_chebyshev_series():
    """Checks coefficients returned by chebyshev_series are correct"""
    for width, terms in [(6, 10), (6, 20)]:
        result = chebyshev_series(torch.tanh, width, terms)

        # check shape
        assert result.shape == torch.Size([terms])

        # check terms
        assert result[0] < 1e-4
        assert torch.isclose(result[-1], torch.tensor(3.5e-2), atol=1e-1) 

def _test_n_k_q_combination(self, n, k, q):
        n_shot_taskloader = DataLoader(self.dataset,
                                       batch_sampler=NShotTaskSampler(self.dataset, 100, n, k, q))

        # Load a single n-shot, k-way task
        for batch in n_shot_taskloader:
            x, y = batch
            break

        # Take just dummy label features and a little bit of noise
        # So distances are never 0
        support = x[:n * k, 1:]
        queries = x[n * k:, 1:]
        support += torch.rand_like(support)
        queries += torch.rand_like(queries)

        distances = pairwise_distances(queries, support, 'cosine')

        # Calculate "attention" as softmax over distances
        attention = (-distances).softmax(dim=1).cuda()

        y_pred = matching_net_predictions(attention, n, k, q)

        self.assertEqual(
            y_pred.shape,
            (q * k, k),
            'Matching Network predictions must have shape (q * k, k).'
        )

        y_pred_sum = y_pred.sum(dim=1)
        self.assertTrue(
            torch.all(
                torch.isclose(y_pred_sum, torch.ones_like(y_pred_sum).double())
            ),
            'Matching Network predictions probabilities must sum to 1 for each '
            'query sample.'
        ) 

def test_tuple_wait(cuda_sleep):
    # In v0.0.3, Wait is applied to only the first tensor on a micro-batch.
    # Under this behavior, if checkpointing was disabled, there's a possibility
    # that gradient accumulations on other tensors are not synchronized
    # properly to the copy stream.
    class Sleep(torch.autograd.Function):
        @staticmethod
        def forward(ctx, x):
            return x.detach()

        @staticmethod
        def backward(ctx, grad):
            with torch.cuda.device(grad.device):
                cuda_sleep(0.05)
            return grad

    class Layer1(nn.Module):
        def forward(self, pair):
            a, b = pair
            return a*1, b*2, b*3

    class Layer2(nn.Module):
        def forward(self, triple):
            a, b, c = triple
            b = Sleep.apply(b)
            return a+b+c

    model = nn.Sequential(Layer1(), Layer2())
    model = GPipe(model, [1, 1], devices=[0, 1], chunks=32, checkpoint='never')

    a = torch.rand(1024, 3, 32, 32, device=0, requires_grad=True)
    b = torch.rand(1024, 3, 32, 32, device=0, requires_grad=True)

    y = model((a, b))
    y.norm().backward()

    torch.cuda.synchronize(0)
    torch.cuda.synchronize(1)

    assert torch.isclose(b.grad.norm().cpu(), torch.tensor(5.000)) 

def test_calculate_precision_float():
    pred = torch.tensor([1, 2, 3, 5, 7, 6], dtype=torch.long)
    true = torch.tensor([1, 2, 3, 4, 5], dtype=torch.long)
    pred = VariableShapeList.from_tensors([pred])
    true = VariableShapeList.from_tensors([true])

    precision = vsl_precision(pred, true)
    assert torch.isclose(precision[0], torch.tensor(4/6)) 

def test_calculate_recall_float():
    pred = torch.tensor([1, 2, 3, 5, 7, 6], dtype=torch.long)
    true = torch.tensor([1, 2, 3, 4, 5], dtype=torch.long)
    pred = VariableShapeList.from_tensors([pred])
    true = VariableShapeList.from_tensors([true])

    recall = vsl_recall(pred, true)
    assert torch.isclose(recall[0], torch.tensor(4/5)) 

def test_chebyshev_series(self):
        """Checks coefficients returned by chebyshev_series are correct"""
        for width, terms in [(6, 10), (6, 20)]:
            result = chebyshev_series(torch.tanh, width, terms)
            # check shape
            self.assertTrue(result.shape == torch.Size([terms]))
            # check terms
            self.assertTrue(result[0] < 1e-4)
            self.assertTrue(torch.isclose(result[-1], torch.tensor(3.5e-2), atol=1e-1)) 

def test_bernoulli(self):
        for size in [(10,), (10, 10), (10, 10, 10)]:
            probs = torch.rand(size)
            randvec = crypten.bernoulli(probs)
            self.assertTrue(randvec.size() == size, "Incorrect size")
            tensor = randvec.get_plain_text()
            self.assertTrue(((tensor == 0) + (tensor == 1)).all(), "Invalid values")

        probs = torch.Tensor(int(1e4)).fill_(0.2)
        randvec = crypten.bernoulli(probs).get_plain_text()
        frac_zero = float((randvec == 0).sum()) / randvec.nelement()
        self.assertTrue(math.isclose(frac_zero, 0.8, rel_tol=1e-1, abs_tol=1e-1)) 

def test_takes_log_with_nllloss(self, net_cls, module_cls, data):
        net = net_cls(module_cls, criterion=nn.NLLLoss, max_epochs=1)
        net.initialize()

        mock_loss = Mock(side_effect=nn.NLLLoss())
        net.criterion_.forward = mock_loss
        net.partial_fit(*data)  # call partial_fit to avoid re-initialization

        # check that loss was called with log-probabilities
        for (y_log, _), _ in mock_loss.call_args_list:
            assert (y_log < 0).all()
            y_proba = torch.exp(y_log)
            assert torch.isclose(torch.ones(len(y_proba)), y_proba.sum(1)).all()

    # classifier-specific test 

def test_takes_no_log_without_nllloss(self, net_cls, module_cls, data):
        net = net_cls(module_cls, criterion=nn.BCELoss, max_epochs=1)
        net.initialize()

        mock_loss = Mock(side_effect=nn.NLLLoss())
        net.criterion_.forward = mock_loss
        net.partial_fit(*data)  # call partial_fit to avoid re-initialization

        # check that loss was called with raw probabilities
        for (y_out, _), _ in mock_loss.call_args_list:
            assert not (y_out < 0).all()
            assert torch.isclose(torch.ones(len(y_out)), y_out.sum(1)).all()

    # classifier-specific test 

def isclose(x, y, rtol=1e-05, atol=1e-08, equal_nan=False):
    """
    Parameters:
    -----------
    x, y : tensor
        Input tensors to compare.
    rtol : float
        The relative tolerance parameter (see Notes).
    atol : float
        The absolute tolerance parameter (see Notes).
    equal_nan : bool
        Whether to compare NaN’s as equal. If True, NaN’s in x will be considered equal to NaN’s in y in the output array.

    Returns:
    --------
    isclose : boolean tensor of where a and b are equal within the given tolerance.
        If both x and y are scalars, returns a single boolean value.
    """
    t1, t2 = __sanitize_close_input(x, y)

    # no sanitation for shapes of x and y needed, torch.isclose raises relevant errors
    _local_isclose = torch.isclose(t1._DNDarray__array, t2._DNDarray__array, rtol, atol, equal_nan)

    # If x is distributed, then y is also distributed along the same axis
    if t1.comm.is_distributed() and t1.split is not None:
        output_gshape = stride_tricks.broadcast_shape(t1.gshape, t2.gshape)
        res = torch.empty(output_gshape).bool()
        t1.comm.Allgather(_local_isclose, res)
        result = factories.array(res, dtype=types.bool, device=t1.device, split=t1.split)
    else:
        if _local_isclose.dim() == 0:
            # both x and y are scalars, return a single boolean value
            result = bool(factories.array(_local_isclose).item())
        else:
            result = factories.array(_local_isclose, dtype=types.bool, device=t1.device)

    return result 

def check_triplets_are_hardest(
    ids_anchor: List[int],
    ids_pos: List[int],
    ids_neg: List[int],
    labels: List[int],
    distmat: Tensor,
) -> None:
    """
    Args:
        ids_anchor: anchor indexes of selected triplets
        ids_pos: positive indexes of selected triplets
        ids_neg: negative indexes of selected triplets
        labels: labels of the samples in the batch
        distmat: distances between features
    """
    ids_all = set(range(len(labels)))

    for i_a, i_p, i_n in zip(ids_anchor, ids_pos, ids_neg):
        ids_label = set(find_value_ids(it=labels, value=labels[i_a]))

        ids_pos_cur = np.array(list(ids_label - {i_a}), int)
        ids_neg_cur = np.array(list(ids_all - ids_label), int)

        assert torch.isclose(
            distmat[i_a, ids_pos_cur].max(), distmat[i_a, i_p]
        )

        assert torch.isclose(
            distmat[i_a, ids_neg_cur].min(), distmat[i_a, i_n]
        ) 

def test_class_weighted_reducer(self):
        class_weights = torch.tensor([1, 0.9, 1, 0.1, 0, 0, 0, 0, 0, 0])
        reducer = ClassWeightedReducer(class_weights)
        batch_size = 100
        num_classes = 10
        embedding_size = 64
        embeddings = torch.randn(batch_size, embedding_size)
        labels = torch.randint(0,num_classes,(batch_size,))
        pair_indices = (torch.randint(0,batch_size,(batch_size,)), torch.randint(0,batch_size,(batch_size,)))
        triplet_indices = pair_indices + (torch.randint(0,batch_size,(batch_size,)),)
        losses = torch.randn(batch_size)

        for indices, reduction_type in [(torch.arange(batch_size), "element"),
                                        (pair_indices, "pos_pair"),
                                        (pair_indices, "neg_pair"),
                                        (triplet_indices, "triplet")]:
            loss_dict = {"loss": {"losses": losses, "indices": indices, "reduction_type": reduction_type}}
            output = reducer(loss_dict, embeddings, labels)
            correct_output = 0
            for i in range(len(losses)):
                if reduction_type == "element":
                    batch_idx = indices[i]
                else:
                    batch_idx = indices[0][i]
                class_label = labels[batch_idx]
                correct_output += losses[i]*class_weights[class_label]
            correct_output /= len(losses)
            self.assertTrue(torch.isclose(output,correct_output)) 

def test_regular_face_regularizer(self):
        temperature = 0.1
        num_classes = 10
        embedding_size = 512
        reg_weight = 0.1
        loss_func = NormalizedSoftmaxLoss(temperature=temperature, 
                                            num_classes=num_classes, 
                                            embedding_size=embedding_size,
                                            regularizer=RegularFaceRegularizer(),
                                            reg_weight=reg_weight)

        embeddings = torch.nn.functional.normalize(torch.randn((180, embedding_size), requires_grad=True, dtype=torch.float))
        labels = torch.randint(low=0, high=10, size=(180,))

        loss = loss_func(embeddings, labels)
        loss.backward()

        weights = torch.nn.functional.normalize(loss_func.W, p=2, dim=0)
        logits = torch.matmul(embeddings, weights)
        correct_class_loss = torch.nn.functional.cross_entropy(logits/temperature, labels)

        weight_cos_matrix = torch.matmul(weights.t(), weights)
        weight_cos_matrix.fill_diagonal_(float('-inf'))
        correct_reg_loss = 0
        for i in range(num_classes):
            correct_reg_loss += torch.max(weight_cos_matrix[i])
        correct_reg_loss /= num_classes

        correct_total_loss = correct_class_loss+(correct_reg_loss*reg_weight)
        self.assertTrue(torch.isclose(loss, correct_total_loss)) 

def test_center_invariant_regularizer(self):
        temperature = 0.1
        num_classes = 10
        embedding_size = 512
        reg_weight = 0.1
        loss_func = NormalizedSoftmaxLoss(temperature=temperature, 
                                            num_classes=num_classes, 
                                            embedding_size=embedding_size,
                                            regularizer=CenterInvariantRegularizer(),
                                            reg_weight=reg_weight)

        embeddings = torch.nn.functional.normalize(torch.randn((180, embedding_size), requires_grad=True, dtype=torch.float))
        labels = torch.randint(low=0, high=10, size=(180,))

        loss = loss_func(embeddings, labels)
        loss.backward()

        weights = torch.nn.functional.normalize(loss_func.W, p=2, dim=0)
        logits = torch.matmul(embeddings, weights)
        correct_class_loss = torch.nn.functional.cross_entropy(logits/temperature, labels)

        correct_reg_loss = 0
        average_squared_weight_norms = 0
        for i in range(num_classes):
            average_squared_weight_norms += torch.norm(loss_func.W[:,i], p=2)**2
        average_squared_weight_norms /= num_classes
        for i in range(num_classes):
            deviation = torch.norm(loss_func.W[:,i], p=2)**2 - average_squared_weight_norms
            correct_reg_loss += (deviation**2) / 4
        correct_reg_loss /= num_classes

        correct_total_loss = correct_class_loss+(correct_reg_loss*reg_weight)
        self.assertTrue(torch.isclose(loss, correct_total_loss)) 

def test_multi_similarity_loss(self):
        alpha, beta, base = 0.1, 40, 0.5
        loss_func = MultiSimilarityLoss(alpha=alpha, beta=beta, base=base)

        embedding_angles = [0, 20, 40, 60, 80]
        embeddings = torch.tensor([c_f.angle_to_coord(a) for a in embedding_angles], requires_grad=True, dtype=torch.float) #2D embeddings
        labels = torch.LongTensor([0, 0, 1, 1, 2])

        loss = loss_func(embeddings, labels)
        loss.backward()

        pos_pairs = [(0,1), (1,0), (2,3), (3,2)]
        neg_pairs = [(0,2), (0,3), (0,4), (1,2), (1,3), (1,4), (2,0), (2,1), (2,4), (3,0), (3,1), (3,4), (4,0), (4,1), (4,2), (4,3)]

        correct_total = 0
        for i in range(len(embeddings)):
            correct_pos_loss = 0
            correct_neg_loss = 0
            for a,p in pos_pairs:
                if a == i:
                    anchor, positive = embeddings[a], embeddings[p]
                    correct_pos_loss += torch.exp(-alpha*(torch.matmul(anchor,positive)-base))
            if correct_pos_loss > 0:
                correct_pos_loss = (1/alpha) * torch.log(1+correct_pos_loss)

            for a,n in neg_pairs:
                if a == i:
                    anchor, negative = embeddings[a], embeddings[n]
                    correct_neg_loss += torch.exp(beta*(torch.matmul(anchor,negative)-base))
            if correct_neg_loss > 0:
                correct_neg_loss = (1/beta) * torch.log(1+correct_neg_loss)
            correct_total += correct_pos_loss + correct_neg_loss

        correct_total /= embeddings.size(0)
        self.assertTrue(torch.isclose(loss, correct_total))