InfiniCore/python/infinicore/tensor.py at 7b7eb9906c0b9f3e76c98f76add3f3093527a327 · gongchensu/InfiniCore · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
import ctypes

import _infinicore


class Tensor:
    def __init__(self, tensor):
        """An internal method. Please do not use this directly."""

        self._underlying = tensor

    @property
    def shape(self):
        return self._underlying.shape

    @property
    def dtype(self):
        return self._underlying.dtype

    @property
    def device(self):
        return self._underlying.device

    @property
    def ndim(self):
        return self._underlying.device

    def data_ptr(self):
        return self._underlying.data_ptr

    def size(self, dim=None):
        if dim is None:
            return self.shape

        return self.shape[dim]

    def stride(self, dim=None):
        if dim is None:
            return self._underlying.strides

        return self._underlying.strides[dim]

    def numel(self):
        return self._underlying.numel()

    def is_contiguous(self):
        return self._underlying.is_contiguous()

    def is_is_pinned(self):
        return self._underlying.is_is_pinned()

    def copy_(self, src):
        return Tensor(self._underlying.copy_(src._underlying))

    def to(self, *args, **kwargs):
        return Tensor(
            self._underlying.to(*tuple(arg._underlying for arg in args), **kwargs)
        )

    def as_strided(self, size, stride):
        return Tensor(self._underlying.as_strided(size, stride))

    def contiguous(self):
        return Tensor(self._underlying.contiguous())

    def permute(self, dims):
        return Tensor(self._underlying.permute(dims))

    def view(self, shape):
        return Tensor(self._underlying.view(shape))

    def debug(self, filename=None):
        """Print tensor data or save to file for debugging

        Args:
            filename: Optional filename to save raw binary data. If None, prints to stdout.
        """
        if filename is None:
            self._underlying.debug()
        else:
            self._underlying.debug(filename)


def empty(size, *, dtype=None, device=None, pin_memory=False):
    return Tensor(
        _infinicore.empty(size, dtype._underlying, device._underlying, pin_memory)
    )


def zeros(size, *, dtype=None, device=None, pin_memory=False):
    return Tensor(
        _infinicore.zeros(size, dtype._underlying, device._underlying, pin_memory)
    )


def ones(size, *, dtype=None, device=None, pin_memory=False):
    return Tensor(
        _infinicore.ones(size, dtype._underlying, device._underlying, pin_memory)
    )


def from_blob(data_ptr, size, *, dtype=None, device=None):
    return Tensor(
        _infinicore.from_blob(data_ptr, size, dtype._underlying, device._underlying)
    )