* init commit

* finish writing

* add to docs

* fix docs

* fix typo

* delete new line

* rename to tensor properties

---------

Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com>
This commit is contained in:
Eitan Turok
2024-08-07 10:38:51 -04:00
committed by GitHub
parent 39dda3d042
commit 39c8c9c00a
2 changed files with 126 additions and 7 deletions

View File

@@ -111,3 +111,12 @@
::: tinygrad.Tensor.half
::: tinygrad.Tensor.int
::: tinygrad.Tensor.bool
## Tensor Properties
::: tinygrad.Tensor.ndim
::: tinygrad.Tensor.numel
::: tinygrad.Tensor.element_size
::: tinygrad.Tensor.nbytes
::: tinygrad.Tensor.is_floating_point
::: tinygrad.Tensor.size

View File

@@ -2949,15 +2949,125 @@ class Tensor:
smoothing = label_smoothing * (log_probs.mean(-1) * loss_mask).sum()
return -((1 - label_smoothing) * (log_probs * y).sum() + smoothing) / loss_mask.sum()
# ***** convenience stuff *****
# ***** Tensor Properties *****
@property
def ndim(self) -> int: return len(self.shape)
def numel(self) -> sint: return prod(self.shape)
def element_size(self) -> int: return self.dtype.itemsize
def nbytes(self) -> int: return self.numel() * self.element_size()
def is_floating_point(self) -> bool: return dtypes.is_float(self.dtype)
def size(self, dim=None) -> Union[sint, Tuple[sint, ...]]: return self.shape if dim is None else self.shape[dim]
def ndim(self) -> int:
"""
Returns the number of dimensions in the tensor.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([9])
print(t.ndim)
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[1, 2], [3, 4]])
print(t.ndim)
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
print(t.ndim)
```
"""
return len(self.shape)
def numel(self) -> sint:
"""
Returns the total number of elements in the tensor.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([9])
print(t.numel())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[1, 2], [3, 4]])
print(t.numel())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
print(t.numel())
```
"""
return prod(self.shape)
def element_size(self) -> int:
"""
Returns the size in bytes of an individual element in the tensor.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([True], dtype=dtypes.bool)
print(t.element_size())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([5], dtype=dtypes.int16)
print(t.element_size())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([5.5], dtype=dtypes.float64)
print(t.element_size())
```
"""
return self.dtype.itemsize
def nbytes(self) -> int:
"""
Returns the total number of bytes of all elements in the tensor.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([8, 9], dtype=dtypes.int16)
print(t.nbytes())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([8, 9, 10, 11], dtype=dtypes.int16)
print(t.nbytes())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([8, 9], dtype=dtypes.float64)
print(t.nbytes())
```
"""
return self.numel() * self.element_size()
def is_floating_point(self) -> bool:
"""
Returns `True` if the tensor contains floating point types, i.e. is one of `dtype.half`, `dtype.float`,
`dtype.double`, `dtype.default_float`, `dtype.float16`, `dtype.float32`, `dtype.float64`, `dtype.bfloat16`.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([8.0, 9.0], dtype=dtypes.float)
print(t.is_floating_point())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([8, 9], dtype=dtypes.float)
print(t.is_floating_point())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([8, 9], dtype=dtypes.int)
print(t.is_floating_point())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([8, 9], dtype=dtypes.float64)
print(t.is_floating_point())
```
"""
return dtypes.is_float(self.dtype)
def size(self, dim=None) -> Union[sint, Tuple[sint, ...]]:
"""
Return the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[4, 5, 6], [7, 8, 9]])
print(t.size())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.size(0))
```
```python exec="true" source="above" session="tensor" result="python"
print(t.size(1))
```
"""
return self.shape if dim is None else self.shape[dim]
# ***** cast ops *****