InfiniCore/python/infinicore/tensor.py at c7054b7acfdabee4bee52bc2d6a7c81be7078d93 · 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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import infinicore.device
import infinicore.dtype
from infinicore.lib import _infinicore


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

        self._underlying = underlying

        self._dtype = infinicore.dtype(self._underlying.dtype)

        self._device = infinicore.device._from_infinicore_device(
            self._underlying.device
        )

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

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

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

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

    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):
        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 strided_empty(size, strides, *, dtype=None, device=None, pin_memory=False):
    return Tensor(
        _infinicore.strided_empty(
            size, strides, 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)
    )


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