Python torch.lt() Examples

def regenerate_cache(self):
        """
        Resamples the big matrix and resets the counter of the total
        number of elements in the returned masks.
        """
        low_size = int(self.resolution * self.max_size)
        low_pattern = self.rng.uniform(0, 1, size=(low_size, low_size)) * 255
        low_pattern = torch.from_numpy(low_pattern.astype('float32'))
        pattern = transforms.Compose([
                        transforms.ToPILImage(),
                        transforms.Resize(self.max_size, Image.BICUBIC),
                        transforms.ToTensor(),
        ])(low_pattern[None])[0]
        pattern = torch.lt(pattern, self.density).byte()
        self.pattern = pattern.byte()
        self.points_used = 0 

def reward_ddpg_B(solution, use_cuda):
    """
    Count number of (nonconsecutively) correctly sorted starting from beginning
    
    Tends to converge to [0,2,4,6,8,1,3,5,7,9]

    solution is FloatTensor of dim [batch,n]
    """
    (batch_size, n, m) = solution.size()
    n_correctly_sorted = Variable(torch.zeros(batch_size, 1), requires_grad=False)
    
    if use_cuda:
        n_correctly_sorted = n_correctly_sorted.cuda()

    for i in range(1, m):
        res = torch.lt(solution[:,:,i-1], solution[:,:,i])
        n_correctly_sorted += res.float()

    return torch.div(n_correctly_sorted, m - 1) 

def test_train(self):
        self._metric.train()
        calls = [[torch.FloatTensor([0.0]), torch.LongTensor([0])],
                 [torch.FloatTensor([0.0, 0.1, 0.2, 0.3]), torch.LongTensor([0, 1, 2, 3])]]
        for i in range(len(self._states)):
            self._metric.process(self._states[i])
        self.assertEqual(2, len(self._metric_function.call_args_list))
        for i in range(len(self._metric_function.call_args_list)):
            self.assertTrue(torch.eq(self._metric_function.call_args_list[i][0][0], calls[i][0]).all)
            self.assertTrue(torch.lt(torch.abs(torch.add(self._metric_function.call_args_list[i][0][1], -calls[i][1])), 1e-12).all)
        self._metric_function.reset_mock()
        self._metric.process_final({})

        self.assertEqual(self._metric_function.call_count, 1)
        self.assertTrue(torch.eq(self._metric_function.call_args_list[0][0][1], torch.LongTensor([0, 1, 2, 3, 4])).all)
        self.assertTrue(torch.lt(torch.abs(torch.add(self._metric_function.call_args_list[0][0][0], -torch.FloatTensor([0.0, 0.1, 0.2, 0.3, 0.4]))), 1e-12).all) 

def reward_ddpg_A(solution, use_cuda):
    """
    Count number of consecutively correctly sorted for each sample in minibatch
    starting from beginning
    
    Very hard to randomly find enough non-zero samples - reward is very sparse!

    solution is FloatTensor of dim [batch,n]
    """
    (batch_size, n, m) = solution.size()
    n_correctly_sorted = Variable(torch.zeros(batch_size, 1), requires_grad=False)
    mask = Variable(torch.ones(batch_size, 1).byte(), requires_grad=False)
    if use_cuda:
        n_correctly_sorted = n_correctly_sorted.cuda()
        mask = mask.cuda()

    for i in range(1, m):
        res = torch.lt(solution[:,:,i-1], solution[:,:,i])
        mask.data &= res.data
        n_correctly_sorted[mask] += 1

    return torch.div(n_correctly_sorted, m - 1) 

def forward(self, words, frequent_tuning=False):
        if frequent_tuning and self.training:

            padding_mask = words.eq(0).long()

            # Fine-tuning - N the most frequent
            fine_tune_mask = torch.lt(words, self.threshold_index) * padding_mask.eq(
                0
            )  # < threshold_index
            fine_tune_words = words * fine_tune_mask.long()

            fine_tune_embedded = self.fine_tune_word_embedding(fine_tune_words)
            fine_tune_embedded = f.masked_zero(fine_tune_embedded, fine_tune_mask)

            # Fixed - under N frequent
            fixed_mask = torch.ge(words, self.threshold_index)  # >= threshold_index

            fixed_embedeed = self.fixed_word_embedding(words).detach()  # Fixed
            fixed_embedeed = f.masked_zero(fixed_embedeed, fixed_mask)

            embedded_words = fine_tune_embedded + fixed_embedeed
        else:
            embedded_words = self.fixed_word_embedding(words)

        return self.dropout(embedded_words) 

def check_monitor_top_k(self, current):
        less_than_k_models = len(self.best_k_models) < self.save_top_k
        if less_than_k_models:
            return True

        if not isinstance(current, torch.Tensor):
            rank_zero_warn(
                f'{current} is supposed to be a `torch.Tensor`. Saving checkpoint may not work correctly.'
                f' HINT: check the value of {self.monitor} in your validation loop', RuntimeWarning
            )
            current = torch.tensor(current)

        monitor_op = {
            "min": torch.lt,
            "max": torch.gt,
        }[self.mode]

        return monitor_op(current, self.best_k_models[self.kth_best_model_path]) 

def get_morph(batch):

    #Not very nice but we do not have access to value comming from opt.gpuid command line parameter here.
    use_cuda = batch.src[0].is_cuda

    # morph_index = batch.morph.data.transpose(0, 1)  # [ seqLen x batch_size ] ==> [ batch_size x seqLen ]

    # morph_voc = batch.dataset.fields['morph'].vocab.stoi

    morph_index = batch.morph.view((batch.src[0].data.size()[0], 6, batch.src[0].data.size()[1]))
    morph_index = morph_index.permute(2, 0, 1).contiguous()



    # morph_index = torch.LongTensor(morph_index)
    morph_mask = torch.lt(torch.eq(morph_index, 1), 1).float()
    # morph_index = autograd.Variable(morph_index)
    # morph_mask = autograd.Variable(torch.FloatTensor(morph_mask), requires_grad=False)
    if use_cuda:
        morph_index = morph_index.cuda()
        morph_mask = morph_mask.cuda()

    return morph_index, morph_mask 

def intersectionAndUnion(batch_data, pred, numClass):
    (imgs, segs, infos) = batch_data
    _, preds = torch.max(pred.data.cpu(), dim=1)

    # compute area intersection
    intersect = preds.clone()
    intersect[torch.ne(preds, segs)] = -1

    area_intersect = torch.histc(intersect.float(),
                                 bins=numClass,
                                 min=0,
                                 max=numClass - 1)

    # compute area union:
    preds[torch.lt(segs, 0)] = -1
    area_pred = torch.histc(preds.float(),
                            bins=numClass,
                            min=0,
                            max=numClass - 1)
    area_lab = torch.histc(segs.float(),
                           bins=numClass,
                           min=0,
                           max=numClass - 1)
    area_union = area_pred + area_lab - area_intersect
    return area_intersect, area_union 

def _jit_linear_cg_updates(
    result, alpha, residual_inner_prod, eps, beta, residual, precond_residual, mul_storage, is_zero, curr_conjugate_vec
):
    # # Update result
    # # result_{k} = result_{k-1} + alpha_{k} p_vec_{k-1}
    result = torch.addcmul(result, alpha, curr_conjugate_vec, out=result)

    # beta_{k} = (precon_residual{k}^T r_vec_{k}) / (precon_residual{k-1}^T r_vec_{k-1})
    beta.resize_as_(residual_inner_prod).copy_(residual_inner_prod)
    torch.mul(residual, precond_residual, out=mul_storage)
    torch.sum(mul_storage, -2, keepdim=True, out=residual_inner_prod)

    # Do a safe division here
    torch.lt(beta, eps, out=is_zero)
    beta.masked_fill_(is_zero, 1)
    torch.div(residual_inner_prod, beta, out=beta)
    beta.masked_fill_(is_zero, 0)

    # Update curr_conjugate_vec
    # curr_conjugate_vec_{k} = precon_residual{k} + beta_{k} curr_conjugate_vec_{k-1}
    curr_conjugate_vec.mul_(beta).add_(precond_residual) 

def wrapper_gmask(opt):
    # batchsize should be 1 for mask_global
    mask_global = torch.ByteTensor(1, 1, \
                                        opt.fineSize, opt.fineSize)

    res = 0.06  # the lower it is, the more continuous the output will be. 0.01 is too small and 0.1 is too large
    density = 0.25
    MAX_SIZE = 350
    maxPartition = 30
    low_pattern = torch.rand(1, 1, int(res * MAX_SIZE), int(res * MAX_SIZE)).mul(255)
    pattern = F.interpolate(low_pattern, (MAX_SIZE, MAX_SIZE), mode='bilinear').detach()
    low_pattern = None
    pattern.div_(255)
    pattern = torch.lt(pattern, density).byte()  # 25% 1s and 75% 0s
    pattern = torch.squeeze(pattern).byte()

    gMask_opts = {}
    gMask_opts['pattern'] = pattern
    gMask_opts['MAX_SIZE'] = MAX_SIZE
    gMask_opts['fineSize'] = opt.fineSize
    gMask_opts['maxPartition'] = maxPartition
    gMask_opts['mask_global'] = mask_global
    return create_gMask(gMask_opts)  # create an initial random mask. 

def test_terner_connect_sto_forward():
    x = torch.Tensor([1,0,0.45,-1,-0.9]).view(1,-1)

    results = list()
    for i in range(1000):
        temp_result = TernaryConnectStochastic.apply(x)
        # Tensor must have only -1 , 0 , 1 values
        assert not torch.any(torch.lt(torch.abs(temp_result-1),1e-8)*torch.lt(torch.abs(temp_result),1e-8))
        results.append(temp_result) 

    result = torch.cat(results,0 )
    result = torch.sum(result, 0)/1000
    
    assert equals(
        result,
        torch.Tensor([1,0,0.45,-1,-0.9]).view(1,-1),
        5e-2) 

def maxOfTwo(x, y):
    z = x.clone()
    maskYLarger = torch.lt(x, y)
    z[maskYLarger.detach()] = y[maskYLarger.detach()]
    return z 

def assertTensorAlmostEqual(self, expected, actual):
        diff = torch.all(
            torch.lt(torch.abs(torch.add(expected, -actual)), 1e-4))
        if not diff:
            self.fail("Tensors didn't match but were supposed to {} vs"
                      " {}".format(expected, actual)) 

def lt(t1, t2):
    """
    Element-wise rich less than comparison between values from operand t1 with respect to values of
    operand t2 (i.e. t1 < t2), not commutative.
    Takes the first and second operand (scalar or tensor) whose elements are to be compared as argument.

    Parameters
    ----------
    t1: tensor or scalar
        The first operand to be compared less than second operand

    t2: tensor or scalar
        The second operand to be compared greater than first operand

    Returns
    -------
    result: ht.DNDarray
        A uint8-tensor holding 1 for all elements in which values of t1 are less than values of t2,
        0 for all other elements

    Examples
    -------
    >>> import heat as ht
    >>> T1 = ht.float32([[1, 2],[3, 4]])
    >>> ht.lt(T1, 3.0)
    tensor([[1, 1],
            [0, 0]], dtype=torch.uint8)
    >>> T2 = ht.float32([[2, 2], [2, 2]])
    >>> ht.lt(T1, T2)
    tensor([[1, 0],
            [0, 0]], dtype=torch.uint8)
    """
    return operations.__binary_op(torch.lt, t1, t2) 

def compute(self, left, right) -> torch.Tensor:
        return torch.lt(left, right) 

def forward(self, indices_):
		"""
		Arguments:
			indices_ (LongTensor: any shape): indices matrix where each index points to a word vector.
		Return:
			embedded (FloatTensor: any shape x embedding_dim): each index is replaced by a word vector.
		"""
		if self.use_cuda:
			indices = indices_.view(-1).to(self.devices[0])

			embedded = torch.empty(indices.size(0), self.embedding_dim).to(self.devices[0])
			idxs = torch.arange(indices.size(0)).to(self.devices[0], dtype=torch.long)

			for i in range(self.num_pages):
				mask_i = torch.min(torch.ge(indices, i * self.page_size), torch.lt(indices, (i+1) * self.page_size))
				mask_idx = torch.masked_select(idxs, mask_i)
				if mask_idx.dim() == 0:
					continue

				indices_i = torch.index_select(indices, 0, mask_idx) - i * self.page_size
				indices_i = indices_i.to(self.devices[i])

				try:
					value_i = self.embeddings[i](indices_i).to(self.devices[0])
				except Exception:
					print("LargeEmbedding - %s, %s" % (indices_i, i * self.page_size))
					print("LargeEmbedding - %s" % self.devices[i])
					print("LargeEmbedding - %s" % self.embeddings[i])
					print("LargeEmbedding - %s" % indices_i.get_device()) if self.use_cuda else None
				# embedded.index_copy_(0, mask_idx, value_i)
				embedded[mask_idx, :] = value_i
			dim = list(indices_.size()) + [self.embedding_dim]
			embedded = embedded.view(*dim)
		else:
			embedded = self.embeddings[0](indices_)

		return embedded 

def __init__(self, monitor: str = 'val_loss', min_delta: float = 0.0, patience: int = 3,
                 verbose: bool = False, mode: str = 'auto', strict: bool = True):
        super().__init__()
        self.monitor = monitor
        self.patience = patience
        self.verbose = verbose
        self.strict = strict
        self.min_delta = min_delta
        self.wait_count = 0
        self.stopped_epoch = 0
        self.mode = mode

        if mode not in self.mode_dict:
            if self.verbose > 0:
                log.info(f'EarlyStopping mode {mode} is unknown, fallback to auto mode.')
            self.mode = 'auto'

        if self.mode == 'auto':
            if self.monitor == 'acc':
                self.mode = 'max'
            else:
                self.mode = 'min'
            if self.verbose > 0:
                log.info(f'EarlyStopping mode set to {self.mode} for monitoring {self.monitor}.')

        self.min_delta *= 1 if self.monitor_op == torch.gt else -1
        self.best_score = torch_inf if self.monitor_op == torch.lt else -torch_inf 

def __lt__(self, other):
        if not isinstance(other, TypeWrapper):
            return False
        if isinstance(self.value, Tensor) and isinstance(other.value, Tensor):
            return torch.lt(self.value, other.value).prod().item()
        elif isinstance(self.value, np.ndarray) and isinstance(other.value, np.ndarray):
            return np.less(self.value, other.value).prod()
        else:
            return self.value < other.value 

def reward_ddpg_C(solution, use_cuda, nco=False):
    """
    Max size substring of correctly sorted numbers

    Exploration gets very tricky near global optimum..
    e.g., [0, 1, 2, 4, 5, 6, 7, 9, 8, 3] ??
    --> swap(9,3)
    --> swap(7,3)
    --> swap(6,3)
    --> swap(5,3)
    --> swap(4,3)
    solution is FloatTensor of dim [1,n]
    """
    (batch_size, n, m) = solution.size()
    
    if not nco:
        longest = Variable(torch.ones(batch_size, 1), requires_grad=False)
        current = Variable(torch.ones(batch_size, 1), requires_grad=False)
    else:
        longest = Variable(torch.ones(batch_size), requires_grad=False)
        current = Variable(torch.ones(batch_size), requires_grad=False)
        
    if use_cuda:
        longest = longest.cuda()
        current = current.cuda()
    for i in range(1, m):
        res = torch.lt(solution[:,:, i-1], solution[:,:,i])
        current += res.float()
        current[torch.eq(res, 0)] = 1
        mask = torch.gt(current, longest)
        longest[mask] = current[mask]

    return torch.div(longest, m) 

def forward(self, x, routing_weights):
        x_orig = x
        B, C, H, W = x.shape
        weight = torch.matmul(routing_weights, self.weight)     # (Expert x out x in x 3x3) --> (B x out x in x 3x3)
        new_weight_shape = (B * self.out_channels, self.in_channels // self.groups) + self.kernel_size
        weight = weight.view(new_weight_shape)                  # (B*out x in x 3 x 3)
        bias = None
        if self.bias is not None:
            bias = torch.matmul(routing_weights, self.bias)
            bias = bias.view(B * self.out_channels)
        # move batch elements with channels so each batch element can be efficiently convolved with separate kernel
        x = x.view(1, B * C, H, W)
        if self.dynamic_padding:
            out = conv2d_same(
                x, weight, bias, stride=self.stride, padding=self.padding,
                dilation=self.dilation, groups=self.groups * B)
        else:
            out = F.conv2d(
                x, weight, bias, stride=self.stride, padding=self.padding,
                dilation=self.dilation, groups=self.groups * B)

        # out : (1 x B*out x ...)
        out = out.permute([1, 0, 2, 3]).view(B, self.out_channels, out.shape[-2], out.shape[-1])

        # out2 = self.forward_legacy(x_orig, routing_weights)
        # lt = torch.lt(torch.abs(torch.add(out, -out2)), 1e-8)
        # assert torch.all(lt), torch.abs(torch.add(out, -out2))[lt]
        # print('checked')
        return out 

def forward(self, estDisp, gtDisp):

        if not torch.is_tensor(gtDisp):
            raise TypeError('ground truth disparity map is expected to be tensor, got {}'.format(type(gtDisp)))
        if not torch.is_tensor(estDisp):
            raise TypeError('estimated disparity map is expected to be tensor, got {}'.format(type(estDisp)))

        assert estDisp.shape == gtDisp.shape

        if gtDisp.dim() == 2:  # single image H x W
            h, w = gtDisp.size(0), gtDisp.size(1)
            gtDisp = gtDisp.view(1, 1, h, w)
            estDisp = estDisp.view(1, 1, h, w)

        if gtDisp.dim() == 3:  # multi image B x H x W
            b, h, w = gtDisp.size(0), gtDisp.size(1), gtDisp.size(2)
            gtDisp = gtDisp.view(b, 1, h, w)
            estDisp = estDisp.view(b, 1, h, w)

        if gtDisp.dim() == 4:
            if gtDisp.size(1) == 1:  # mult image B x 1 x H x W
                self.gtDisp = gtDisp
                self.estDisp = estDisp
            else:
                raise ValueError('2nd dimension size should be 1, got {}'.format(gtDisp.size(1)))

        confidence_gt_label = torch.lt(torch.abs(self.estDisp - self.gtDisp), self.theta).type_as(self.gtDisp)

        return confidence_gt_label 

def getProb(self):
        # |d - d{gt}| < variance, [BatchSize, maxDisp, Height, Width]
        probability = torch.lt(torch.abs(self.disp_sample - self.gtDisp), self.variance).type_as(self.gtDisp)

        return probability 

def test_validate(self):
        self._metric.eval()
        for i in range(len(self._states)):
            self._metric.process(self._states[i])
        self._metric_function.assert_not_called()
        self._metric.process_final_validate({})

        self.assertEqual(self._metric_function.call_count, 1)
        self.assertTrue(torch.eq(self._metric_function.call_args_list[0][0][1], torch.LongTensor([0, 1, 2, 3, 4])).all)
        self.assertTrue(torch.lt(torch.abs(torch.add(self._metric_function.call_args_list[0][0][0], -torch.FloatTensor([0.0, 0.1, 0.2, 0.3, 0.4]))), 1e-12).all) 

def log1mexp(U, eps=1e-3):
    """
    Compute log(1 - exp(u)) for u <= 0.
    """
    res = torch.log1p(-torch.exp(U))

    # |U| << 1 requires care for numerical stability:
    # 1 - exp(U) = -U + o(U)
    small = torch.lt(U.abs(), eps)
    res[small] = torch.log(-U[small])

    return res 

def prune_vanilla_elementwise(param, sparsity, fn_importance=lambda x: x.abs()):
    """
    element-wise vanilla pruning
    :param param: torch.(cuda.)Tensor, weight of conv/fc layer
    :param sparsity: float, pruning sparsity
    :param fn_importance: function, inputs 'param' and returns the importance of
                                    each position in 'param',
                                    default=lambda x: x.abs()
    :return:
        torch.(cuda.)ByteTensor, mask for zeros
    """
    sparsity = min(max(0.0, sparsity), 1.0)
    if sparsity == 1.0:
        return torch.zeros_like(param).byte()
    num_el = param.numel()
    importance = fn_importance(param)
    num_pruned = int(math.ceil(num_el * sparsity))
    num_stayed = num_el - num_pruned
    if sparsity <= 0.5:
        _, topk_indices = torch.topk(importance.view(num_el), k=num_pruned,
                                     dim=0, largest=False, sorted=False)
        mask = torch.zeros_like(param).byte()
        param.view(num_el).index_fill_(0, topk_indices, 0)
        mask.view(num_el).index_fill_(0, topk_indices, 1)
    else:
        thr = torch.min(torch.topk(importance.view(num_el), k=num_stayed,
                                   dim=0, largest=True, sorted=False)[0])
        mask = torch.lt(importance, thr)
        param.masked_fill_(mask, 0)
    return mask 

def forward(self, pos_context, pos_path, neg_context, neg_path):
        pos_context_embedding = torch.sum(self.word_embeddings(pos_context), dim=1, keepdim=True)
        pos_path_embedding = self.context_embeddings(pos_path)
        pos_score = torch.bmm(pos_context_embedding, pos_path_embedding.transpose(2, 1)).squeeze()
        neg_context_embedding = torch.sum(self.word_embeddings(neg_context), dim=1, keepdim=True)
        neg_path_embedding = self.context_embeddings(neg_path)
        neg_score = torch.bmm(neg_context_embedding, neg_path_embedding.transpose(2, 1)).squeeze()
        pos_sigmoid_score = torch.lt(torch.sigmoid(pos_score), 0.5)
        neg_sigmoid_score = torch.gt(torch.sigmoid(neg_score), 0.5)
        sigmoid_score = torch.cat((pos_sigmoid_score, neg_sigmoid_score))
        sigmoid_score = torch.sum(sigmoid_score, dim=0).item() / sigmoid_score.size(0)
        return sigmoid_score 

def forward(self, pos_target, pos_path, neg_target, neg_path):
        pos_target_embedding = torch.sum(self.word_embeddings(pos_target), dim=1, keepdim=True)
        pos_path_embedding = self.context_embeddings(pos_path)
        pos_score = torch.bmm(pos_target_embedding, pos_path_embedding.transpose(2, 1)).squeeze()
        neg_target_embedding = torch.sum(self.word_embeddings(neg_target), dim=1, keepdim=True)
        neg_path_embedding = self.context_embeddings(neg_path)
        neg_score = torch.bmm(neg_target_embedding, neg_path_embedding.transpose(2, 1)).squeeze()
        pos_sigmoid_score = torch.lt(torch.sigmoid(pos_score), 0.5)
        neg_sigmoid_score = torch.gt(torch.sigmoid(neg_score), 0.5)
        sigmoid_score = torch.cat((pos_sigmoid_score, neg_sigmoid_score))
        sigmoid_score = torch.sum(sigmoid_score, dim=0).item() / sigmoid_score.size(0)
        return sigmoid_score 

def dist_acc(dists, thr=0.5):
    """
    Return percentage below threshold while ignoring values with a -1
    """
    dist_cal = torch.ne(dists, -1)
    num_dist_cal = dist_cal.sum()
    if num_dist_cal > 0:
        return torch.lt(dists[dist_cal], thr).float().sum() / num_dist_cal
    else:
        return -1 

def logits_mask(coords, logits, num_points_per_object):
    """
    Use logits to sample points
    :param coords: coords of points, FloatTensor[B, 3, N]
    :param logits: binary classification logits, FloatTensor[B, 2, N]
    :param num_points_per_object: M, #points per object after masking, int
    :return:
        selected_coords: FloatTensor[B, 3, M]
        masked_coords_mean: mean coords of selected points, FloatTensor[B, 3]
        mask: mask to select points, BoolTensor[B, N]
    """
    batch_size, _, num_points = coords.shape
    mask = torch.lt(logits[:, 0, :], logits[:, 1, :])   # [B, N]
    num_candidates = torch.sum(mask, dim=-1, keepdim=True)  # [B, 1]
    masked_coords = coords * mask.view(batch_size, 1, num_points)  # [B, C, N]
    masked_coords_mean = torch.sum(masked_coords, dim=-1) / torch.max(num_candidates,
                                                                      torch.ones_like(num_candidates)).float()  # [B, C]
    selected_indices = torch.zeros((batch_size, num_points_per_object), device=coords.device, dtype=torch.int32)
    for i in range(batch_size):
        current_mask = mask[i]  # [N]
        current_candidates = current_mask.nonzero().view(-1)
        current_num_candidates = current_candidates.numel()
        if current_num_candidates >= num_points_per_object:
            choices = np.random.choice(current_num_candidates, num_points_per_object, replace=False)
            selected_indices[i] = current_candidates[choices]
        elif current_num_candidates > 0:
            choices = np.concatenate([
                np.arange(current_num_candidates).repeat(num_points_per_object // current_num_candidates),
                np.random.choice(current_num_candidates, num_points_per_object % current_num_candidates, replace=False)
            ])
            np.random.shuffle(choices)
            selected_indices[i] = current_candidates[choices]
    selected_coords = gather(masked_coords - masked_coords_mean.view(batch_size, -1, 1), selected_indices)
    return selected_coords, masked_coords_mean, mask 

def equals(a, b, epsilon=1e-12):
    return torch.all(torch.lt( torch.abs(a-b), epsilon ))