From 39c8c9c00a9521ae160dcaf867eaef78bc5e50af Mon Sep 17 00:00:00 2001
From: Eitan Turok <150733043+eitanturok@users.noreply.github.com>
Date: Wed, 7 Aug 2024 10:38:51 -0400
Subject: [PATCH] Add docs (#5942)

* 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>
---
 docs/tensor/ops.md |   9 ++++
 tinygrad/tensor.py | 124 ++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 126 insertions(+), 7 deletions(-)

diff --git a/docs/tensor/ops.md b/docs/tensor/ops.md
index 81cd8a823c..4e1e2a051b 100644
--- a/docs/tensor/ops.md
+++ b/docs/tensor/ops.md
@@ -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
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 7e15ca64d2..68bc79f5ad 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -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 *****