Python torch.float64() Examples

def sphere_compliment_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    complement = torch.rand(shape[-1], 1, dtype=torch.float64)

    Q, _ = geoopt.linalg.batch_linalg.qr(complement)
    P = -Q @ Q.transpose(-1, -2)
    P[..., torch.arange(P.shape[-2]), torch.arange(P.shape[-2])] += 1

    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = (ex @ P.t()) / torch.norm(ex @ P.t())
    v = (ev - (x @ ev) * x) @ P.t()

    manifold = geoopt.Sphere(complement=complement)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact(complement=complement)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case 

def pack(x):
    """
    Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor.

    Parameter:
    ----------
    x: A tensor of shape [batch_size, 4, 1]

    Return:
    ------
    A tensor of shape [batch_size, 4, 4] after appending.

    """
    zeros43 = torch.zeros(
      (x.shape[0], x.shape[1], 4, 3), dtype=torch.float64).to(x.device)
    ret = torch.cat((zeros43, x), dim=3)
    return ret 

def with_zeros(x):
    """
    Append a [0, 0, 0, 1] tensor to a [3, 4] tensor.

    Parameter:
    ---------
    x: Tensor to be appended.

    Return:
    ------
    Tensor after appending of shape [4,4]

    """
    ones = torch.tensor(
      [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float64
    ).expand(x.shape[0],-1,-1).to(x.device)
    ret = torch.cat((x, ones), dim=1)
    return ret 

def test03(self):
        padder = AutoPadder()
        # 测试tensor的情况
        # 0维
        contents = torch.arange(12)
        r_contents = padder(contents, None, contents.dtype, 0)
        self.assertSequenceEqual(r_contents.tolist(), contents.tolist())
        self.assertTrue(r_contents.dtype==contents.dtype)

        # 0维
        contents = [torch.tensor(1) for _ in range(10)]
        self.assertSequenceEqual(padder(contents, None, torch.int64, 0).tolist(), contents)

        # 1维
        contents = torch.randn(3, 4)
        padder(contents, None, torch.float64, 1)

        # 3维
        contents = [torch.randn(3, 4, 4) for _ in range(5)]
        padder(contents, None, torch.float64, 3) 

def with_zeros(x):
    """
    Append a [0, 0, 0, 1] tensor to a [3, 4] tensor.

    Parameter:
    ---------
    x: Tensor to be appended.

    Return:
    ------
    Tensor after appending of shape [4,4]

    """
    ones = torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float64).to(x.device)
    ret = torch.cat((x, ones), dim=0)
    return ret 

def __init__(self, name, in_channels, out_channels, kernel_size,
                 stride=1, padding=0, dilation=1, bias=True):
        super().__init__(name)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.bias = bias

        weight_shape = [out_channels, in_channels, kernel_size, kernel_size]
        self.weight = tvm.te.placeholder(weight_shape, dtype='float64', name=f'{name}_weight')
        if bias:
            self.bias = tvm.te.placeholder([out_channels, ], dtype='float64', name=f'{name}_bias')
        else:
            self.bias = None 

def test_gpu(gpu_id=[0]):
  if len(gpu_id) > 0 and torch.cuda.is_available():
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id[0])
    device = torch.device('cuda')
  else:
    device = torch.device('cpu')
  print(device)

  pose_size = 72
  beta_size = 10

  np.random.seed(9608)
  pose = torch.from_numpy((np.random.rand(pose_size) - 0.5) * 0.4)\
          .type(torch.float64).to(device)
  betas = torch.from_numpy((np.random.rand(beta_size) - 0.5) * 0.06) \
          .type(torch.float64).to(device)
  trans = torch.from_numpy(np.zeros(3)).type(torch.float64).to(device)
  outmesh_path = './smpl_torch.obj'

  model = SMPLModel(device=device)
  result = model(betas, pose, trans)
  model.write_obj(result, outmesh_path) 

def pack(x):
    """
    Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor.

    Parameter:
    ----------
    x: A tensor of shape [batch_size, 4, 1]

    Return:
    ------
    A tensor of shape [batch_size, 4, 4] after appending.

    """
    zeros43 = torch.zeros((x.shape[0], 4, 3), dtype=torch.float64).to(x.device)
    ret = torch.cat((zeros43, x), dim=2)
    return ret 

def torch_dtype_to_np_dtype(dtype):
    dtype_dict = {
            torch.bool    : np.dtype(np.bool),
            torch.uint8   : np.dtype(np.uint8),
            torch.int8    : np.dtype(np.int8),
            torch.int16   : np.dtype(np.int16),
            torch.short   : np.dtype(np.int16),
            torch.int32   : np.dtype(np.int32),
            torch.int     : np.dtype(np.int32),
            torch.int64   : np.dtype(np.int64),
            torch.long    : np.dtype(np.int64),
            torch.float16 : np.dtype(np.float16),
            torch.half    : np.dtype(np.float16),
            torch.float32 : np.dtype(np.float32),
            torch.float   : np.dtype(np.float32),
            torch.float64 : np.dtype(np.float64),
            torch.double  : np.dtype(np.float64),
            }
    return dtype_dict[dtype]


# ---------------------- InferenceEngine internal types ------------------------ 

def sample_action(self, obs):
        """Sample an action for the given observation.
        
        Parameters
        ----------
        obs: A numpy array of shape [obs_dim].
        
        Returns
        -------
        An integer, the action sampled.
        """
        if np.random.random() < self.epsilon:
            return np.random.randint(self.ac_dim)
        if len(obs.shape) != 1 or obs.shape[0] != self.obs_dim:
            raise ValueError(
                "Expected input observation shape [obs_dim], got %s" % str(obs.shape)
            )
        obs = torch.tensor(obs.reshape(1, -1), dtype=torch.float64)
        # When sampling use eval mode.
        return (
            torch.distributions.Categorical(logits=self.eval().double().forward(obs))
            .sample()
            .item()
        ) 

def __init__(self, device=None, model_path='./model.pkl'):
    
    super(SMPLModel, self).__init__()
    with open(model_path, 'rb') as f:
      params = pickle.load(f)
    self.J_regressor = torch.from_numpy(
      np.array(params['J_regressor'].todense())
    ).type(torch.float64)
    self.weights = torch.from_numpy(params['weights']).type(torch.float64)
    self.posedirs = torch.from_numpy(params['posedirs']).type(torch.float64)
    self.v_template = torch.from_numpy(params['v_template']).type(torch.float64)
    self.shapedirs = torch.from_numpy(params['shapedirs']).type(torch.float64)
    self.kintree_table = params['kintree_table']
    self.faces = params['f']
    self.device = device if device is not None else torch.device('cpu')
    for name in ['J_regressor', 'weights', 'posedirs', 'v_template', 'shapedirs']:
      _tensor = getattr(self, name)
      setattr(self, name, _tensor.to(device)) 

def symsqrt(a, cond=None, return_rank=False, dtype=torch.float32):
    """Symmetric square root of a positive semi-definite matrix.
    See https://github.com/pytorch/pytorch/issues/25481"""

    s, u = torch.symeig(a, eigenvectors=True)
    cond_dict = {torch.float32: 1e3 * 1.1920929e-07, torch.float64: 1E6 * 2.220446049250313e-16}

    if cond in [None, -1]:
        cond = cond_dict[dtype]

    above_cutoff = (abs(s) > cond * torch.max(abs(s)))

    psigma_diag = torch.sqrt(s[above_cutoff])
    u = u[:, above_cutoff]

    B = u @ torch.diag(psigma_diag) @ u.t()
    if return_rank:
        return B, len(psigma_diag)
    else:
        return B 

def test_n_additions_via_scalar_multiplication(n, a, dtype, negative, manifold, strict):
    n = torch.as_tensor(n, dtype=a.dtype).requires_grad_()
    y = torch.zeros_like(a)
    for _ in range(int(n.item())):
        y = manifold.mobius_add(a, y)
    ny = manifold.mobius_scalar_mul(n, a)
    if negative:
        tolerance = {
            torch.float32: dict(atol=4e-5, rtol=1e-3),
            torch.float64: dict(atol=1e-5, rtol=1e-3),
        }
    else:
        tolerance = {
            torch.float32: dict(atol=2e-6, rtol=1e-3),
            torch.float64: dict(atol=1e-5, rtol=1e-3),
        }
    tolerant_allclose_check(y, ny, strict=strict, **tolerance[dtype])
    ny.sum().backward()
    assert torch.isfinite(n.grad).all()
    assert torch.isfinite(a.grad).all()
    assert torch.isfinite(manifold.k.grad).all() 

def test_scalar_multiplication_distributive(a, r1, r2, manifold, dtype):
    res = manifold.mobius_scalar_mul(r1 + r2, a)
    res1 = manifold.mobius_add(
        manifold.mobius_scalar_mul(r1, a), manifold.mobius_scalar_mul(r2, a),
    )
    res2 = manifold.mobius_add(
        manifold.mobius_scalar_mul(r1, a), manifold.mobius_scalar_mul(r2, a),
    )
    tolerance = {
        torch.float32: dict(atol=5e-6, rtol=1e-4),
        torch.float64: dict(atol=1e-7, rtol=1e-4),
    }
    np.testing.assert_allclose(res1.detach(), res.detach(), **tolerance[dtype])
    np.testing.assert_allclose(res2.detach(), res.detach(), **tolerance[dtype])
    res.sum().backward()
    assert torch.isfinite(a.grad).all()
    assert torch.isfinite(r1.grad).all()
    assert torch.isfinite(r2.grad).all()
    assert torch.isfinite(manifold.k.grad).all() 

def test_geodesic_segment_length_property(a, b, manifold, dtype):
    extra_dims = len(a.shape)
    segments = 12
    t = torch.linspace(0, 1, segments + 1, dtype=dtype).view(
        (segments + 1,) + (1,) * extra_dims
    )
    gamma_ab_t = manifold.geodesic(t, a, b)
    gamma_ab_t0 = gamma_ab_t[:-1]
    gamma_ab_t1 = gamma_ab_t[1:]
    dist_ab_t0mt1 = manifold.dist(gamma_ab_t0, gamma_ab_t1, keepdim=True)
    speed = manifold.dist(a, b, keepdim=True).unsqueeze(0).expand_as(dist_ab_t0mt1)
    # we have exactly 12 line segments
    tolerance = {
        torch.float32: dict(rtol=1e-5, atol=5e-3),
        torch.float64: dict(rtol=1e-5, atol=5e-3),
    }
    length = speed / segments
    np.testing.assert_allclose(
        dist_ab_t0mt1.detach(), length.detach(), **tolerance[dtype]
    )
    (length + dist_ab_t0mt1).sum().backward()
    assert torch.isfinite(a.grad).all()
    assert torch.isfinite(b.grad).all()
    assert torch.isfinite(manifold.k.grad).all() 

def test_geodesic_segement_unit_property(a, b, manifold, dtype):
    extra_dims = len(a.shape)
    segments = 12
    t = torch.linspace(0, 1, segments + 1, dtype=dtype).view(
        (segments + 1,) + (1,) * extra_dims
    )
    gamma_ab_t = manifold.geodesic_unit(t, a, b)
    gamma_ab_t0 = gamma_ab_t[:1]
    gamma_ab_t1 = gamma_ab_t
    dist_ab_t0mt1 = manifold.dist(gamma_ab_t0, gamma_ab_t1, keepdim=True)
    true_distance_travelled = t.expand_as(dist_ab_t0mt1)
    # we have exactly 12 line segments
    tolerance = {
        torch.float32: dict(atol=2e-4, rtol=5e-5),
        torch.float64: dict(atol=1e-10),
    }
    np.testing.assert_allclose(
        dist_ab_t0mt1.detach(), true_distance_travelled.detach(), **tolerance[dtype]
    )
    (true_distance_travelled + dist_ab_t0mt1).sum().backward()
    assert torch.isfinite(a.grad).all()
    assert torch.isfinite(b.grad).all()
    assert torch.isfinite(manifold.k.grad).all() 

def euclidean_stiefel_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.EuclideanStiefel]
    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    u, _, v = torch.svd(ex)
    x = u @ v.t()
    nonsym = x.t() @ ev
    v = ev - x @ (nonsym + nonsym.t()) / 2

    manifold = geoopt.manifolds.EuclideanStiefel()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.manifolds.EuclideanStiefelExact()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case 

def birkhoff_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.BirkhoffPolytope]
    ex = torch.randn(*shape, dtype=torch.float64).abs()
    ev = torch.randn(*shape, dtype=torch.float64)
    max_iter = 100
    eps = 1e-12
    tol = 1e-5
    iter = 0
    c = 1.0 / (torch.sum(ex, dim=-2, keepdim=True) + eps)
    r = 1.0 / (torch.matmul(ex, c.transpose(-1, -2)) + eps)
    while iter < max_iter:
        iter += 1
        cinv = torch.matmul(r.transpose(-1, -2), ex)
        if torch.max(torch.abs(cinv * c - 1)) <= tol:
            break
        c = 1.0 / (cinv + eps)
        r = 1.0 / ((ex @ c.transpose(-1, -2)) + eps)
    x = ex * (r @ c)

    v = proju_original(x, ev)
    manifold = geoopt.manifolds.BirkhoffPolytope()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case 

def sphere_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = ex / torch.norm(ex)
    v = ev - (x @ ev) * x

    manifold = geoopt.Sphere()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact()
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case 

def poincare_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.PoincareBall]
    ex = torch.randn(*shape, dtype=torch.float64) / 3
    ev = torch.randn(*shape, dtype=torch.float64) / 3
    x = torch.tanh(torch.norm(ex)) * ex / torch.norm(ex)
    ex = x.clone()
    v = ev.clone()
    manifold = geoopt.PoincareBall().to(dtype=torch.float64)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.PoincareBallExact().to(dtype=torch.float64)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case 

def sphere_projection_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.SphereProjection]
    ex = torch.randn(*shape, dtype=torch.float64) / 3
    ev = torch.randn(*shape, dtype=torch.float64) / 3
    x = ex  # default curvature = 0
    ex = x.clone()
    v = ev.clone()
    manifold = geoopt.manifolds.SphereProjection().to(dtype=torch.float64)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.manifolds.SphereProjectionExact().to(dtype=torch.float64)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case 

def sphere_subspace_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    subspace = torch.rand(shape[-1], 2, dtype=torch.float64)

    Q, _ = geoopt.linalg.batch_linalg.qr(subspace)
    P = Q @ Q.t()

    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = (ex @ P.t()) / torch.norm(ex @ P.t())
    v = (ev - (x @ ev) * x) @ P.t()

    manifold = geoopt.Sphere(intersection=subspace)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact(intersection=subspace)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case 

def synchronize_between_processes(self):
        """
        Warning: does not synchronize the deque!
        """
        if not is_dist_avail_and_initialized():
            return
        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
        dist.barrier()
        dist.all_reduce(t)
        t = t.tolist()
        self.count = int(t[0])
        self.total = t[1] 

def __init__(self, name, in_features, out_features, bias=True):
        super().__init__(name)
        self.in_features = in_features
        self.out_features = out_features
        # https://docs.tvm.ai/api/python/topi.html#topi.nn.dense
        # - weight (tvm.te.Tensor) – 2-D with shape [out_dim, in_dim]
        self.weight = tvm.te.placeholder([out_features, in_features], dtype='float64', name=f'{name}_weight')
        if bias:
            self.bias = tvm.te.placeholder([out_features, ], dtype='float64', name=f'{name}_bias')
        else:
            self.bias = None 

def test_sparsemax_grad():
    x = torch.randn(4, 6, dtype=torch.float64, requires_grad=True)
    gradcheck(sparsemax_bisect, (x,), eps=1e-5) 

def test_getitem_v1(self):
        fa = FieldArray("y", [[1.1, 2.2, 3.3, 4.4, 5.5], [1.0, 2.0, 3.0, 4.0, 5.0]], is_input=True)
        self.assertEqual(fa[0], [1.1, 2.2, 3.3, 4.4, 5.5])
        ans = fa[[0, 1]]
        self.assertTrue(isinstance(ans, np.ndarray))
        self.assertTrue(isinstance(ans[0], np.ndarray))
        self.assertEqual(ans[0].tolist(), [1.1, 2.2, 3.3, 4.4, 5.5])
        self.assertEqual(ans[1].tolist(), [1, 2, 3, 4, 5])
        self.assertEqual(ans.dtype, np.float64) 

def test_entmax_grad(alpha):
    alpha = torch.tensor(alpha, dtype=torch.float64, requires_grad=True)
    x = torch.randn(4, 6, dtype=torch.float64, requires_grad=True)
    gradcheck(entmax_bisect, (x, alpha), eps=1e-5) 

def test_entmax_correct_multiple_alphas():
    n = 4
    x = torch.randn(n, 6, dtype=torch.float64, requires_grad=True)
    alpha = 1.05 + torch.rand((n, 1), dtype=torch.float64, requires_grad=True)

    p1 = entmax_bisect(x, alpha)
    p2_ = [
        entmax_bisect(x[i].unsqueeze(0), alpha[i].item()).squeeze()
        for i in range(n)
    ]
    p2 = torch.stack(p2_)

    assert torch.allclose(p1, p2) 

def test_entmax_grad_multiple_alphas():

    n = 4
    x = torch.randn(n, 6, dtype=torch.float64, requires_grad=True)
    alpha = 1.05 + torch.rand((n, 1), dtype=torch.float64, requires_grad=True)
    gradcheck(entmax_bisect, (x, alpha), eps=1e-5) 

def test_support_np_array(self):
        fa = FieldArray("y", np.array([[1.1, 2.2, 3.3, 4.4, 5.5]]), is_input=True)
        self.assertEqual(fa.dtype, np.float64)

        fa.append(np.array([1.1, 2.2, 3.3, 4.4, 5.5]))
        self.assertEqual(fa.dtype, np.float64)

        fa = FieldArray("my_field", np.random.rand(3, 5), is_input=True)
        # in this case, pytype is actually a float. We do not care about it.
        self.assertEqual(fa.dtype, np.float64)