Python torch.DoubleTensor() Examples

def test_dbbox2delta(self):
        """
        encoding format similar to RRPN, except the angle was restricted to [0, 2 pi], dangle was restricted to [0, 1]

        Must Test corner cases
        :return:
        """
        boxlist1 = torch.DoubleTensor([[1, 1, 10, 5, 0],
                                  [1, 1, 10, 5, np.pi/10],
                                  [1, 1, 10, 5, 0],
                                  [30, 100, 60, 34, np.pi/2]
                                  ])

        boxlist2 = torch.DoubleTensor([[1, 1, 5, 8, np.pi/16],
                                  [1, 1, 5, 8, np.pi/16 + np.pi/10],
                                  [1, 1, 10, 5, 0],
                                  [30, 90, 12, 45, np.pi/10]
                                  ])
        expected_targets = torch.DoubleTensor([[0.0000,  0.0000, -0.6931,  0.4700,  0.0312],
                                        [0.0000,  0.0000, -0.6931,  0.4700,  0.0313],
                                        [0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
                                        [-0.1667,  0.0000, -1.6094,  0.2803, 0.8]])

        output = dbbox2delta(boxlist1, boxlist2)
        np.testing.assert_almost_equal(expected_targets.numpy(), output.numpy(), decimal=4) 

def test_dbbox_flip(self):
        """

        :return:
        """
        dbboxes = torch.DoubleTensor([[50, 60, 50, 30, np.pi/3, 23.5, 60, 50, 30, np.pi/3 + np.pi/10],
                                      [30, 20, 30, 30, np.pi/6, 10.33, 20, 30, 30, np.pi/6 - np.pi/11]])
        img_shape = (1024, 1024)

        expected_targets = torch.DoubleTensor([[973, 60, 50, 30, np.pi/3*2, 999.5, 60, 50, 30, np.pi*17/30],
                                               [993, 20, 30, 30, np.pi/6*5, 1012.67, 20, 30, 30, np.pi*61/66]])

        output = dbbox_flip(dbboxes, img_shape)

        np.testing.assert_almost_equal(expected_targets.numpy(), output.numpy(), decimal=6)

    # def dbbox_mapping(self):
    #
    #     dbboxes = torch.DoubleTensor([[50, 60, 50, 30, np.pi/3, 23.5, 60, 50, 30, np.pi/3 + np.pi/10],
    #                                   [30, 20, 30, 30, np.pi/6, 10.33, 20, 30, 30, np.pi/6 - np.pi/11]])
    #     img_shape 

def __init__(self, n, Qpenalty, qp_solver, trueInit=False):
        super().__init__()

        self.qp_solver = qp_solver

        nx = (n**2)**3
        self.Q = Variable(Qpenalty*torch.eye(nx).double().cuda())
        self.Q_idx = spa.csc_matrix(self.Q.detach().cpu().numpy()).nonzero()

        self.G = Variable(-torch.eye(nx).double().cuda())
        self.h = Variable(torch.zeros(nx).double().cuda())
        t = get_sudoku_matrix(n)

        if trueInit:
            self.A = Parameter(torch.DoubleTensor(t).cuda())
        else:
            self.A = Parameter(torch.rand(t.shape).double().cuda())
        self.log_z0 = Parameter(torch.zeros(nx).double().cuda())
        # self.b = Variable(torch.ones(self.A.size(0)).double().cuda())

        if self.qp_solver == 'osqpth':
            t = torch.cat((self.A, self.G), dim=0)
            self.AG_idx = spa.csc_matrix(t.detach().cpu().numpy()).nonzero()

    # @profile 

def test_dbbox2roi(self):

        dbbox_list = [torch.DoubleTensor([[2, 3, 39, 30, np.pi/2],
                                          [3.2, 3, 30, 20, np.pi/3]]),
                      torch.DoubleTensor([[1, 3, 39, 30, np.pi / 2],
                                          [5.2, 3, 30, 20, np.pi / 3]]) ]

        expected_targets = np.array([[0, 2, 3, 39, 30, np.pi/2],
                                     [0, 3.2, 3, 30, 20, np.pi/3],
                                     [1, 1, 3, 39, 30, np.pi / 2],
                                     [1, 5.2, 3, 30, 20, np.pi / 3]
                                     ])

        outputs = dbbox2roi(dbbox_list)

        np.testing.assert_almost_equal(expected_targets, outputs.numpy()) 

def test_droi2dbbox(self):

        drois = np.array([[0, 2, 3, 39, 30, np.pi/2],
                                     [0, 3.2, 3, 30, 20, np.pi/3],
                                     [1, 1, 3, 39, 30, np.pi / 2],
                                     [1, 5.2, 3, 30, 20, np.pi / 3]
                                     ])
        drois = torch.from_numpy(drois)

        outputs = droi2dbbox(drois)

        expected_targets = [torch.DoubleTensor([[2, 3, 39, 30, np.pi/2],
                                          [3.2, 3, 30, 20, np.pi/3]]),
                      torch.DoubleTensor([[1, 3, 39, 30, np.pi / 2],
                                          [5.2, 3, 30, 20, np.pi / 3]])]
        np.testing.assert_equal(len(outputs), 2)
        np.testing.assert_equal(expected_targets[0].shape == outputs[0].shape, True)
        np.testing.assert_equal(expected_targets[1].shape == outputs[1].shape, True)

        # np.testing.assert_almost_equal(expected_targets, outputs.numpy()) 

def convert_robot_one_hot(self, labels):
        """
        Converts a list of robot_id labels to a Torch tensor of one_hot labels
    
        :param labels: List of robot_id labels
        :type labels: list
    
        :returns: A torch tensor of one_hot labels
        :rtype: torch Tensor
        """
        batch_size = len(labels)
        labels_tensor = Variable(torch.LongTensor(labels))
        y_onehot = Variable(torch.DoubleTensor(batch_size, n_rob))
        y_onehot.zero_()
        labels_onehot = y_onehot.scatter_(1, labels_tensor.view(-1, 1), 1)
        return labels_onehot 

def convert_one_hot(self, labels):
        """
        Converts a list of grasp angle labels to a Torch tensor of one_hot labels
    
        :param labels: List of angle labels
        :type labels: list
    
        :returns: A torch tensor of one_hot labels
        :rtype: torch Tensor
        """
        batch_size = len(labels)
        labels_tensor = Variable(torch.LongTensor(labels))
        y_onehot = Variable(torch.DoubleTensor(batch_size, n_class))
        y_onehot.zero_()
        labels_onehot = y_onehot.scatter_(1, labels_tensor.view(-1, 1), 1)
        return labels_onehot 

def calc_relative_scale(skeleton, ref_bone_lengths, joint_tree) -> (float, float):
    """Calculate the factor by which the reference is larger than the query skeleton.

    Args:
        skeleton (torch.DoubleTensor): The query skeleton.
        ref_bone_lengths (torch.DoubleTensor): The reference skeleton bone lengths.
        joint_tree (list of int):

    Returns:
        The average scale factor.
    """

    bone_lengths = cartesian_to_spherical(
        absolute_to_parent_relative(ensure_cartesian(skeleton, d=3), joint_tree)
    )[:, 0]

    non_zero = bone_lengths.gt(1e-6)
    if non_zero.sum() == 0: return 0
    ratio = (ref_bone_lengths / bone_lengths).masked_select(non_zero)

    return ratio.median().item() 

def test_sum_metric_inputs(self):

        passing_inputs_and_outputs = [
            (2, 2.0),
            (-5.0, -5.0),
            (torch.LongTensor([[-1]]), -1.0),
            (torch.DoubleTensor([34.68]), 34.68),
        ]
        for input_, output in passing_inputs_and_outputs:
            actual_output = SumMetric(input_).value()
            self.assertEqual(actual_output, output)

        failing_inputs = [
            ('4', AssertionError),
            ([6.8], AssertionError),
            (torch.Tensor([1, 3.8]), ValueError),  # Tensor has more than 1 element
        ]
        for input_, error in failing_inputs:
            with self.assertRaises(error):
                SumMetric(input_) 

def class_accuracy(num_classes, output_transform=lambda x: x, device=None):
    """Calculates class accuracy

    Args:
        num_classes (int): number of classes
        output_transform (callable, optional): a callable that is used to transform the
            output into the form expected by the metric.

    Returns:
        MetricsLambda

    """
    cm = ignite.metrics.ConfusionMatrix(num_classes=num_classes, output_transform=output_transform, device=device)
    # Increase floating point precision and pass to CPU
    cm = cm.type(torch.DoubleTensor)

    acc_cls = cm.diag() / (cm.sum(dim=1) + 1e-15)

    return acc_cls 

def pixelwise_accuracy(num_classes, output_transform=lambda x: x, device=None):
    """Calculates class accuracy

    Args:
        num_classes (int): number of classes
        output_transform (callable, optional): a callable that is used to transform the
            output into the form expected by the metric.

    Returns:
        MetricsLambda

    """
    cm = ignite.metrics.ConfusionMatrix(num_classes=num_classes, output_transform=output_transform, device=device)
    # Increase floating point precision and pass to CPU
    cm = cm.type(torch.DoubleTensor)

    pix_cls = ignite.metrics.confusion_matrix.cmAccuracy(cm)

    return pix_cls 

def test_horovod_allreduce_grad_average(self):
        """Test the correctness of the allreduce averaged gradient."""
        hvd.init()
        dtypes = [torch.IntTensor, torch.LongTensor,
                  torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.IntTensor, torch.cuda.LongTensor,
                       torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            tensor = tensor.type(dtype)
            tensor = torch.autograd.Variable(tensor, requires_grad=True)
            summed = hvd.allreduce(tensor, average=True)

            summed.backward(torch.ones([17] * dim))
            grad_out = tensor.grad.data.numpy()

            expected = np.ones([17] * dim)
            err = np.linalg.norm(expected - grad_out)
            self.assertLess(err, 0.00000001,
                            "gradient %s differs from expected %s, "
                            "error: %s" % (grad_out, expected, str(err))) 

def __init__(self, dataset, indices=None, num_samples=None, callback_get_label=None):
                
        # if indices is not provided, 
        # all elements in the dataset will be considered
        self.indices = list(range(len(dataset))) \
            if indices is None else indices

        # define custom callback
        self.callback_get_label = callback_get_label

        # if num_samples is not provided, 
        # draw `len(indices)` samples in each iteration
        self.num_samples = len(self.indices) \
            if num_samples is None else num_samples
            
        # distribution of classes in the dataset 
        label_to_count = {}
        for idx in self.indices:
            label = self._get_label(dataset, idx)
            if label in label_to_count:
                label_to_count[label] += 1
            else:
                label_to_count[label] = 1
                
        # weight for each sample
        weights = [1.0 / label_to_count[self._get_label(dataset, idx)]
                   for idx in self.indices]
        self.weights = torch.DoubleTensor(weights) 

def __init__(self, n, Qpenalty, trueInit=False):
        super().__init__()
        nx = (n**2)**3
        self.nx = nx

        spTensor = torch.cuda.sparse.DoubleTensor
        iTensor = torch.cuda.LongTensor
        dTensor = torch.cuda.DoubleTensor

        self.Qi = iTensor([range(nx), range(nx)])
        self.Qv = Variable(dTensor(nx).fill_(Qpenalty))
        self.Qsz = torch.Size([nx, nx])

        self.Gi = iTensor([range(nx), range(nx)])
        self.Gv = Variable(dTensor(nx).fill_(-1.0))
        self.Gsz = torch.Size([nx, nx])
        self.h = Variable(torch.zeros(nx).double().cuda())

        t = get_sudoku_matrix(n)
        neq = t.shape[0]
        if trueInit:
            I = t != 0
            self.Av = Parameter(dTensor(t[I]))
            Ai_np = np.nonzero(t)
            self.Ai = torch.stack((torch.LongTensor(Ai_np[0]),
                                   torch.LongTensor(Ai_np[1]))).cuda()
            self.Asz = torch.Size([neq, nx])
        else:
            # TODO: This is very dense:
            self.Ai = torch.stack((iTensor(list(range(neq))).unsqueeze(1).repeat(1, nx).view(-1),
                                iTensor(list(range(nx))).repeat(neq)))
            self.Av = Parameter(dTensor(neq*nx).uniform_())
            self.Asz = torch.Size([neq, nx])
        self.b = Variable(torch.ones(neq).double().cuda()) 

def sparse_solve_kkt_inverse(H_, A_, C_tilde, rx, rs, rz, ry, ns):
    nineq, nz, neq, nBatch = ns

    if neq > 0:
        g_ = torch.cat([rx, rs], 1).squeeze(0).numpy()
        h_ = torch.cat([rz, ry], 1).squeeze(0).numpy()
    else:
        g_ = torch.cat([rx, rs], 1).squeeze(0).numpy()
        h_ = rz.squeeze(0).numpy()

    full_mat = bmat([[H_, A_.transpose()],
                     [A_, C_tilde]], format='csc')
    full_res = np.concatenate([g_, h_], 0)
    sol = splu(full_mat).solve(full_res)
    # sol = spsolve(full_mat, full_res)

    dx = sol[:nz]
    ds = sol[nz:nz+nineq]
    dz = sol[nz+nineq:nz+nineq+nineq]
    dy = sol[nz+nineq+nineq:] if neq > 0 else None

    dx = torch.DoubleTensor(dx).unsqueeze(0)
    ds = torch.DoubleTensor(ds).unsqueeze(0)
    dz = torch.DoubleTensor(dz).unsqueeze(0)
    dy = torch.DoubleTensor(dy).unsqueeze(0) if neq > 0 else None

    return dx, ds, dz, dy 

def update(self, output):
        y_pred, y = self._check_shape(output)
        self._check_type((y_pred, y))

        if self._type == "binary":
            y_pred = y_pred.view(-1)
            y = y.view(-1)
        elif self._type == "multiclass":
            num_classes = y_pred.size(1)
            y = to_onehot(y.view(-1), num_classes=num_classes)
            indices = torch.max(y_pred, dim=1)[1].view(-1)
            y_pred = to_onehot(indices, num_classes=num_classes)
        elif self._type == "multilabel":
            # if y, y_pred shape is (N, C, ...) -> (C, N x ...)
            num_classes = y_pred.size(1)
            y_pred = torch.transpose(y_pred, 1, 0).reshape(num_classes, -1)
            y = torch.transpose(y, 1, 0).reshape(num_classes, -1)

        y = y.type_as(y_pred)
        correct = y * y_pred
        all_positives = y_pred.sum(dim=0).type(torch.DoubleTensor)  # Convert from int cuda/cpu to double cpu

        if correct.sum() == 0:
            true_positives = torch.zeros_like(all_positives)
        else:
            true_positives = correct.sum(dim=0)
        # Convert from int cuda/cpu to double cpu
        # We need double precision for the division true_positives / all_positives
        true_positives = true_positives.type(torch.DoubleTensor)

        if self._type == "multilabel":
            self._true_positives = torch.cat([self._true_positives, true_positives], dim=0)
            self._positives = torch.cat([self._positives, all_positives], dim=0)
        else:
            self._true_positives += true_positives
            self._positives += all_positives 

def default_collate(batch):
    "Puts each data field into a tensor with outer dimension batch size"

    error_msg = "batch must contain tensors, numbers, dicts or lists; found {}"
    elem_type = type(batch[0])
    if torch.is_tensor(batch[0]):
        out = None
        if _use_shared_memory:
            # If we're in a background process, concatenate directly into a
            # shared memory tensor to avoid an extra copy
            numel = sum([x.numel() for x in batch])
            storage = batch[0].storage()._new_shared(numel)
            out = batch[0].new(storage)
        return torch.stack(batch, 0, out=out)
    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
            and elem_type.__name__ != 'string_':
        elem = batch[0]
        if elem_type.__name__ == 'ndarray':
            # array of string classes and object
            if re.search('[SaUO]', elem.dtype.str) is not None:
                raise TypeError(error_msg.format(elem.dtype))

            return torch.stack([torch.from_numpy(b) for b in batch], 0)
        if elem.shape == ():  # scalars
            py_type = float if elem.dtype.name.startswith('float') else int
            return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))
    elif isinstance(batch[0], int):
        return torch.LongTensor(batch)
    elif isinstance(batch[0], float):
        return torch.DoubleTensor(batch)
    elif isinstance(batch[0], string_classes):
        return batch
    elif isinstance(batch[0], collections.Mapping):
        return {key: default_collate([d[key] for d in batch]) for key in batch[0]}
    elif isinstance(batch[0], collections.Sequence):
        transposed = zip(*batch)
        return [default_collate(samples) for samples in transposed]

    raise TypeError((error_msg.format(type(batch[0])))) 

def update(self, output):
        y_pred, y = self._check_shape(output)
        self._check_type((y_pred, y))

        if self._type == "binary":
            y_pred = y_pred.view(-1)
            y = y.view(-1)
        elif self._type == "multiclass":
            num_classes = y_pred.size(1)
            y = to_onehot(y.view(-1), num_classes=num_classes)
            indices = torch.max(y_pred, dim=1)[1].view(-1)
            y_pred = to_onehot(indices, num_classes=num_classes)
        elif self._type == "multilabel":
            # if y, y_pred shape is (N, C, ...) -> (C, N x ...)
            num_classes = y_pred.size(1)
            y_pred = torch.transpose(y_pred, 1, 0).reshape(num_classes, -1)
            y = torch.transpose(y, 1, 0).reshape(num_classes, -1)

        y = y.type_as(y_pred)
        correct = y * y_pred
        actual_positives = y.sum(dim=0).type(torch.DoubleTensor)  # Convert from int cuda/cpu to double cpu

        if correct.sum() == 0:
            true_positives = torch.zeros_like(actual_positives)
        else:
            true_positives = correct.sum(dim=0)

        # Convert from int cuda/cpu to double cpu
        # We need double precision for the division true_positives / actual_positives
        true_positives = true_positives.type(torch.DoubleTensor)

        if self._type == "multilabel":
            self._true_positives = torch.cat([self._true_positives, true_positives], dim=0)
            self._positives = torch.cat([self._positives, actual_positives], dim=0)
        else:
            self._true_positives += true_positives
            self._positives += actual_positives 

def reset(self):
        self._true_positives = torch.DoubleTensor(0) if self._is_multilabel else 0
        self._positives = torch.DoubleTensor(0) if self._is_multilabel else 0 

def auc(actual, expected, included_joints=None):
    # This range of thresholds mimics `mpii_compute_3d_pck.m`, which is provided as part of the
    # MPI-INF-3DHP test data release.
    thresholds = torch.linspace(0, 150, 31).tolist()

    pck_values = torch.DoubleTensor(len(thresholds))
    for i, threshold in enumerate(thresholds):
        pck_values[i] = pck(actual, expected, included_joints, threshold=threshold)
    return pck_values.mean().item() 

def __init__(self, weights, factor=2.):
        self.weights = torch.DoubleTensor(weights)
        assert factor >= 1.
        self.num_samples = int(len(self.weights) * factor) 

def default_collate(batch):
    "Puts each data field into a tensor with outer dimension batch size"

    error_msg = "batch must contain tensors, numbers, dicts or lists; found {}"
    elem_type = type(batch[0])
    if torch.is_tensor(batch[0]):
        out = None
        if _use_shared_memory:
            # If we're in a background process, concatenate directly into a
            # shared memory tensor to avoid an extra copy
            numel = sum([x.numel() for x in batch])
            storage = batch[0].storage()._new_shared(numel)
            out = batch[0].new(storage)
        return torch.stack(batch, 0, out=out)
    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
            and elem_type.__name__ != 'string_':
        elem = batch[0]
        if elem_type.__name__ == 'ndarray':
            # array of string classes and object
            if re.search('[SaUO]', elem.dtype.str) is not None:
                raise TypeError(error_msg.format(elem.dtype))

            return torch.stack([torch.from_numpy(b) for b in batch], 0)
        if elem.shape == ():  # scalars
            py_type = float if elem.dtype.name.startswith('float') else int
            return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))
    elif isinstance(batch[0], int):
        return torch.LongTensor(batch)
    elif isinstance(batch[0], float):
        return torch.DoubleTensor(batch)
    elif isinstance(batch[0], string_classes):
        return batch
    elif isinstance(batch[0], collections.Mapping):
        return {key: default_collate([d[key] for d in batch]) for key in batch[0]}
    elif isinstance(batch[0], collections.Sequence):
        transposed = zip(*batch)
        return [default_collate(samples) for samples in transposed]

    raise TypeError((error_msg.format(type(batch[0])))) 

def test_horovod_broadcast_grad(self):
        """Test the correctness of the broadcast gradient."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        # This test does not apply if there is only one worker.
        if size == 1:
            return

        dtypes = [torch.ByteTensor, torch.CharTensor, torch.ShortTensor,
                  torch.IntTensor, torch.LongTensor, torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.ByteTensor, torch.cuda.CharTensor, torch.cuda.ShortTensor,
                       torch.cuda.IntTensor, torch.cuda.LongTensor, torch.cuda.FloatTensor,
                       torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        root_ranks = list(range(size))
        for dtype, dim, root_rank in itertools.product(dtypes, dims, root_ranks):
            tensor = torch.FloatTensor(*([17] * dim)).fill_(1).mul_(rank)
            tensor = tensor.type(dtype)
            tensor = torch.autograd.Variable(tensor, requires_grad=True)

            broadcasted_tensor = hvd.broadcast(tensor, root_rank)
            broadcasted_tensor.backward(torch.ones([17] * dim))
            grad_out = tensor.grad.data.numpy()

            c = size if rank == root_rank else 0
            expected = np.ones([17] * dim) * c
            err = np.linalg.norm(expected - grad_out)
            self.assertLess(err, 0.00000001,
                            "gradient %s differs from expected %s, "
                            "error: %s" % (grad_out, expected, str(err))) 

def test_horovod_broadcast(self):
        """Test that the broadcast correctly broadcasts 1D, 2D, 3D tensors."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        # This test does not apply if there is only one worker.
        if size == 1:
            return

        dtypes = [torch.ByteTensor, torch.CharTensor, torch.ShortTensor,
                  torch.IntTensor, torch.LongTensor, torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.ByteTensor, torch.cuda.CharTensor, torch.cuda.ShortTensor,
                       torch.cuda.IntTensor, torch.cuda.LongTensor, torch.cuda.FloatTensor,
                       torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        root_ranks = list(range(size))
        for dtype, dim, root_rank in itertools.product(dtypes, dims, root_ranks):
            tensor = torch.FloatTensor(*([17] * dim)).fill_(1).mul_(rank)
            root_tensor = torch.FloatTensor(*([17] * dim)).fill_(1).mul_(root_rank)
            tensor = tensor.type(dtype)
            root_tensor = root_tensor.type(dtype)
            broadcasted_tensor = hvd.broadcast(tensor, root_rank)
            if rank != root_rank:
                assert (tensor == root_tensor).max() == 0, \
                    'hvd.broadcast modifies source tensor'
            assert (broadcasted_tensor.data == root_tensor).min() == 1, \
                'hvd.broadcast produces incorrect broadcasted tensor' 

def test_horovod_allgather_variable_size(self):
        """Test that the allgather correctly gathers 1D, 2D, 3D tensors,
        even if those tensors have different sizes along the first dim."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        dtypes = [torch.ByteTensor, torch.CharTensor, torch.ShortTensor,
                  torch.IntTensor, torch.LongTensor, torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.ByteTensor, torch.cuda.CharTensor, torch.cuda.ShortTensor,
                       torch.cuda.IntTensor, torch.cuda.LongTensor, torch.cuda.FloatTensor,
                       torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            # Support tests up to MPI Size of 35
            if size > 35:
                break

            tensor_sizes = [17, 32, 81, 12, 15, 23, 22] * 5
            tensor_sizes = tensor_sizes[:size]

            tensor = torch.FloatTensor(
                *([tensor_sizes[rank]] + [17] * (dim - 1))).fill_(1).mul_(rank)
            tensor = tensor.type(dtype)
            gathered = hvd.allgather(tensor)

            expected_size = sum(tensor_sizes)
            assert list(gathered.shape) == [expected_size] + [17] * (dim - 1)

            for i in range(size):
                rank_size = [tensor_sizes[i]] + [17] * (dim - 1)
                rank_tensor = gathered[sum(
                    tensor_sizes[:i]):sum(tensor_sizes[:i + 1])]
                assert list(rank_tensor.shape) == rank_size
                assert rank_tensor.data.min() == i
                assert rank_tensor.data.max() == i 

def test_horovod_allgather(self):
        """Test that the allgather correctly gathers 1D, 2D, 3D tensors."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        dtypes = [torch.ByteTensor, torch.CharTensor, torch.ShortTensor,
                  torch.IntTensor, torch.LongTensor, torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.ByteTensor, torch.cuda.CharTensor, torch.cuda.ShortTensor,
                       torch.cuda.IntTensor, torch.cuda.LongTensor, torch.cuda.FloatTensor,
                       torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            tensor = torch.FloatTensor(*([17] * dim)).fill_(1).mul_(rank)
            tensor = tensor.type(dtype)
            gathered = hvd.allgather(tensor)

            assert list(gathered.shape) == [17 * size] + [17] * (dim - 1)

            for i in range(size):
                rank_tensor = gathered[i * 17:(i + 1) * 17]
                assert list(rank_tensor.shape) == [17] * dim, \
                    'hvd.allgather produces incorrect gathered shape'
                assert rank_tensor.data.min() == i, 'hvd.allgather produces incorrect gathered tensor'
                assert rank_tensor.data.max() == i, 'hvd.allgather produces incorrect gathered tensor' 

def test_horovod_allreduce_multi_gpu(self):
        """Test that the allreduce works on multiple GPUs."""
        # Only do this test if there are GPUs available.
        if not torch.cuda.is_available():
            return

        hvd.init()
        local_rank = hvd.local_rank()
        size = hvd.size()

        iter = 0
        dtypes = [torch.cuda.IntTensor, torch.cuda.LongTensor,
                  torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            iter += 1
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            device = local_rank * 2 + (iter + local_rank) % 2
            tensor = tensor.cuda(device).type(dtype)
            multiplied = tensor * size
            hvd.allreduce_(tensor, average=False)
            max_difference = tensor.sub(multiplied).max()

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [torch.cuda.IntTensor, torch.cuda.LongTensor]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            assert max_difference <= threshold, 'hvd.allreduce produces incorrect results' 

def test_horovod_allreduce_inplace(self):
        """Test that the allreduce correctly sums 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        dtypes = [torch.IntTensor, torch.LongTensor,
                  torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.IntTensor, torch.cuda.LongTensor,
                       torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            tensor = tensor.type(dtype)
            multiplied = tensor * size
            hvd.allreduce_(tensor, average=False)
            max_difference = tensor.sub(multiplied).max()

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [torch.IntTensor, torch.LongTensor,
                                      torch.cuda.IntTensor, torch.cuda.LongTensor]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            assert max_difference <= threshold, 'hvd.allreduce produces incorrect results' 

def test_horovod_allreduce_average(self):
        """Test that the allreduce correctly sums 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        dtypes = [torch.IntTensor, torch.LongTensor,
                  torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.IntTensor, torch.cuda.LongTensor,
                       torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            tensor = tensor.type(dtype)
            averaged = hvd.allreduce(tensor, average=True)
            max_difference = averaged.data.sub(tensor).max()

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [torch.IntTensor, torch.LongTensor,
                                      torch.cuda.IntTensor, torch.cuda.LongTensor]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            assert max_difference <= threshold, 'hvd.allreduce produces incorrect results' 

def test_nms_device_and_dtypes_cpu():
    """
    CommandLine:
        xdoctest -m tests/test_nms.py test_nms_device_and_dtypes_cpu
    """
    iou_thr = 0.6
    base_dets = np.array([[49.1, 32.4, 51.0, 35.9, 0.1],
                          [49.3, 32.9, 51.0, 35.3, 0.05],
                          [35.3, 11.5, 39.9, 14.5, 0.9],
                          [35.2, 11.7, 39.7, 15.7, 0.3]])

    base_expected_suppressed = np.array([[35.3, 11.5, 39.9, 14.5, 0.9],
                                         [49.1, 32.4, 51.0, 35.9, 0.1]])
    # CPU can handle float32 and float64
    dets = base_dets.astype(np.float32)
    expected_suppressed = base_expected_suppressed.astype(np.float32)
    suppressed, inds = nms(dets, iou_thr)
    assert dets.dtype == suppressed.dtype
    assert np.array_equal(suppressed, expected_suppressed)

    dets = torch.FloatTensor(base_dets)
    expected_suppressed = torch.FloatTensor(base_expected_suppressed)
    suppressed, inds = nms(dets, iou_thr)
    assert dets.dtype == suppressed.dtype
    assert torch.equal(suppressed, expected_suppressed)

    dets = base_dets.astype(np.float64)
    expected_suppressed = base_expected_suppressed.astype(np.float64)
    suppressed, inds = nms(dets, iou_thr)
    assert dets.dtype == suppressed.dtype
    assert np.array_equal(suppressed, expected_suppressed)

    dets = torch.DoubleTensor(base_dets)
    expected_suppressed = torch.DoubleTensor(base_expected_suppressed)
    suppressed, inds = nms(dets, iou_thr)
    assert dets.dtype == suppressed.dtype
    assert torch.equal(suppressed, expected_suppressed)