Add further docstrings and comments (#29)

* Add docstrings to tensor_create.py

* Add docstrings

* Add docstrings to index_slice.oy

* Add docstrings and comments to reshape.py

* Add docstrings and comments to reshape.py

* Small improvements and precommit

LICENSE

@@ -4,4 +4,4 @@ Permission is hereby granted, free of charge, to any person obtaining a copy of
 
 The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

edugrad/__init__.py

@@ -1 +1 @@
-from edugrad.tensor import Tensor  # noqa: F401
+from edugrad.tensor import Tensor

edugrad/_tensor/tensor_broadcasted_binary_mlops.py

@@ -1,5 +1,8 @@
-"""This module implements broadcasted binary operations for Tensors, providing element-wise arithmetic operations that
-support broadcasting for tensors of different shapes."""
+"""Consists broadcasted binary operations for Tensors.
+
+These operations provide element-wise arithmetic operations that support broadcasting for tensors of different shapes.
+
+"""
 from __future__ import annotations
 
 import math
@@ -66,8 +69,7 @@ def _broadcasted(tensor: Tensor, y: Tensor | float, reverse: bool = False) -> tu
 
 
 def _to_float(tensor: Tensor, x: Tensor | float):
-    """Converts a tensor to float32 dtype if it is not already a Tensor and if it is suitable for certain operations
-    where float32 dtype is required.
+    """Converts a tensor to float32 dtype.
 
     Args:
         tensor (Tensor): The reference tensor to check compatibility.
@@ -81,6 +83,7 @@ def _to_float(tensor: Tensor, x: Tensor | float):
 
     return (
         x.data.base.op.arg
+        # tensor is not already a Tensor and suitable for certain operations where float32 dtype is required.
         if isinstance(x, Tensor)
         and x.data.is_unrealized_contiguous_const()
         and not x.requires_grad
@@ -139,7 +142,7 @@ def pow(tensor: Tensor, x: Tensor | float, reverse=False) -> Tensor:
 
     x = tensor._to_float(x)
     if x.__class__ is not Tensor and not reverse:
-        # simple pow identities
+        # Simple pow identities
         if x < 0:
             return tensor.reciprocal().pow(-x)
         if x == 3.0:
@@ -153,24 +156,24 @@ def pow(tensor: Tensor, x: Tensor | float, reverse=False) -> Tensor:
     if not isinstance(x, Tensor) and reverse and x > 0:
         return tensor.mul(math.log(x)).exp()
     ar = tensor.abs().log().mul(x).exp() if not reverse or isinstance(x, Tensor) else tensor.mul(math.log(abs(x))).exp()
-    # correct sign of negative numbers raised to a power (cos has a period of 2pi so we use it here to get the oddness of the power)
+    # Correct sign of negative numbers raised to a power (cos has a period of 2pi so we use it here to get the oddness of the power)
     sign = (
         (x * math.pi).cos()
         if isinstance(x, Tensor)
         else math.cos(x * math.pi)
         if not reverse
         else (tensor * math.pi).cos()
     )
-    # we only need to correct the sign if the base is negative
+    # We only need to correct the sign if the base is negative
     base_sign = (
         (tensor.sign() if not reverse else x.sign() if isinstance(x, Tensor) else math.copysign(1, x)) - 1
     ) / -2
-    # we need 0 to be positive so we need to correct base_sign when the base is 0
+    # We need 0 to be positive so we need to correct base_sign when the base is 0
     base_sign = base_sign - (
         1.5
         * (1 - (tensor.sign().abs() if not reverse else x.sign().abs() if isinstance(x, Tensor) else abs(int(bool(x)))))
     )
-    # inject nan if the base is negative and the power is not an integer
+    # Inject nan if the base is negative and the power is not an integer
     to_nan = (
         ((x - x.trunc()) * 1e10).abs().clip(0, 1)
         if isinstance(x, Tensor)

edugrad/_tensor/tensor_create.py

@@ -1,7 +1,16 @@
+"""Contains low-level operation entry points and helper functions for tensor creation and manipulation.
+
+It includes functions for creating tensors with specific properties (like being empty, random, or having specific
+values) and for random number generation.
+
+"""
+
 from __future__ import annotations
 import time
 import math
 
+from typing import Optional, Any
+
 from edugrad.dtypes import DType, dtypes
 from edugrad.helpers import argfix, prod, shape_int
 from edugrad.data import TensorData
@@ -11,7 +20,20 @@
 # creation low-level op entrypoint *****
 
 
-def _loadop(op, sz, dtype: DType | None = None, arg=None, **kwargs):
+def _loadop(op: LoadOps, sz: int, dtype: DType | None = None, arg: Any = None, **kwargs) -> Tensor:
+    """Internal helper function to create a Tensor with a specified operation.
+
+    Args:
+    - op: Operation to be performed for tensor creation.
+    - sz: Size of the tensor to be created.
+    - dtype: Data type of the tensor. Defaults to Tensor's default type if not provided.
+    - arg: Additional argument for the operation.
+    - kwargs: Additional keyword arguments.
+
+    Returns:
+    - Tensor: A new tensor created with the specified operation.
+
+    """
     from edugrad.tensor import Tensor
 
     assert isinstance(sz, int), f"cannot create with symbolic size {sz}"
@@ -20,7 +42,8 @@ def _loadop(op, sz, dtype: DType | None = None, arg=None, **kwargs):
     )
 
 
-def empty(*shape, **kwargs):
+def empty(*shape, **kwargs) -> Tensor:
+    """Creates an uninitialized tensor with the given shape."""
     from edugrad.tensor import Tensor
 
     return Tensor._loadop(LoadOps.EMPTY, prod(shape := argfix(*shape)), **kwargs).reshape(shape)
@@ -30,12 +53,23 @@ def empty(*shape, **kwargs):
 
 
 def manual_seed(seed=0):
+    """Sets the manual seed for random number generation."""
     from edugrad.tensor import Tensor
 
     Tensor._seed = seed
 
 
-def rand(*shape, **kwargs):
+def rand(*shape, **kwargs) -> Tensor:
+    """Creates a tensor with elements uniformly distributed between 0 and 1.
+
+    Args:
+    - shape: Variable length argument list for the dimensions of the tensor.
+    - kwargs: Additional keyword arguments.
+
+    Returns:
+    - Tensor: A tensor with random elements uniformly distributed.
+
+    """
     from edugrad.tensor import Tensor
 
     Tensor._seed += 1
@@ -46,33 +80,49 @@ def rand(*shape, **kwargs):
 # creation helper functions
 
 
-def full(shape: tuple[shape_int, ...], fill_value, **kwargs):
+def full(shape: tuple[shape_int, ...], fill_value, **kwargs) -> Tensor:
+    """Creates a tensor filled entirely with the specified fill value."""
     from edugrad.tensor import Tensor
 
     return Tensor(fill_value, **kwargs).reshape([1] * len(new_shape := argfix(shape))).expand(new_shape)
 
 
-def zeros(*shape, **kwargs):
+def zeros(*shape, **kwargs) -> Tensor:
+    """Creates a tensor filled entirely with zeros."""
     from edugrad.tensor import Tensor
 
     return Tensor.full(argfix(*shape), 0, **kwargs)
 
 
-def ones(*shape, **kwargs):
+def ones(*shape, **kwargs) -> Tensor:
+    """Creates a tensor filled entirely with ones."""
     from edugrad.tensor import Tensor
 
     return Tensor.full(argfix(*shape), 1, **kwargs)
 
 
-def arange(start, stop, step, **kwargs):
+def arange(start: int | float, stop: int | float | None, step: int | float, **kwargs) -> Tensor:
+    """Creates a 1D tensor with a sequence of numbers from start to stop with a step size.
+
+    Args:
+    - start: The start of the sequence.
+    - stop: The end of the sequence.
+    - step: The step size between each number in the sequence.
+    - kwargs: Additional keyword arguments.
+
+    Returns:
+    - Tensor: A 1D tensor containing a sequence of numbers.
+
+    """
     from edugrad.tensor import Tensor
 
     if stop is None:
         stop, start = start, 0
     return Tensor.full((math.ceil((stop - start) / step),), step, **kwargs).cumsum() + (start - step)
 
 
-def eye(dim: int, **kwargs):
+def eye(dim: int, **kwargs) -> Tensor:
+    """Creates a 2D identity tensor."""
     from edugrad.tensor import Tensor
 
     return (
@@ -84,25 +134,29 @@ def eye(dim: int, **kwargs):
     )
 
 
-def full_like(self, fill_value, **kwargs):
+def full_like(tensor, fill_value, **kwargs) -> Tensor:
+    """Creates a tensor with the same shape as the given tensor, filled with a specified value."""
     from edugrad.tensor import Tensor
 
-    return Tensor.full(self.shape, fill_value=fill_value, dtype=kwargs.pop("dtype", self.dtype), **kwargs)
+    return Tensor.full(tensor.shape, fill_value=fill_value, dtype=kwargs.pop("dtype", tensor.dtype), **kwargs)
 
 
-def zeros_like(self, **kwargs):
-    return self.full_like(0, **kwargs)
+def zeros_like(tensor, **kwargs) -> Tensor:
+    """Creates a tensor with the same shape as the given tensor, filled with zeros."""
+    return tensor.full_like(0, **kwargs)
 
 
-def ones_like(self, **kwargs):
-    return self.full_like(1, **kwargs)
+def ones_like(tensor, **kwargs) -> Tensor:
+    """Creates a tensor with the same shape as the given tensor, filled with ones."""
+    return tensor.full_like(1, **kwargs)
 
 
 # -----------------------------------------------------------------------------------------------------------------------
 # random number generation
 
 
 def randn(*shape, dtype: DType | None, **kwargs) -> Tensor:
+    """Creates a tensor with elements sampled from a standard normal distribution."""
     from edugrad.tensor import Tensor
 
     # https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
@@ -117,25 +171,33 @@ def randn(*shape, dtype: DType | None, **kwargs) -> Tensor:
 
 
 def randint(*shape, low, high, **kwargs) -> Tensor:
+    """Creates a tensor with elements sampled uniformly from the discrete interval [low, high)."""
     from edugrad.tensor import Tensor
 
     return (Tensor.rand(*shape, **kwargs) * (high - low) + low).cast(dtypes.int32)
 
 
 def normal(*shape, mean, std, **kwargs) -> Tensor:
+    """Creates a tensor with elements sampled from a normal (Gaussian) distribution."""
     from edugrad.tensor import Tensor
 
     return (std * Tensor.randn(*shape, **kwargs)) + mean
 
 
 def uniform(*shape, low, high, **kwargs) -> Tensor:
+    """Creates a tensor with elements uniformly distributed over the interval [low, high)."""
     from edugrad.tensor import Tensor
 
     dtype = kwargs.pop("dtype", Tensor.default_type)
     return ((high - low) * Tensor.rand(*shape, **kwargs)).cast(dtype) + low
 
 
 def scaled_uniform(*shape, **kwargs) -> Tensor:
+    """Creates a scaled tensor with elements uniformly distributed over the interval [-1.0, 1.0)
+
+    It is scaled by the inverse square root of the product of the tensor's shape.
+
+    """
     from edugrad.tensor import Tensor
 
     return Tensor.uniform(*shape, low=-1.0, high=1.0, **kwargs).mul(prod(shape) ** -0.5)

edugrad/_tensor/tensor_index_slice.py

@@ -1,4 +1,4 @@
-from typing import Sequence, Optional, Tuple
+from typing import Sequence, Optional, Tuple, Union
 from collections import defaultdict
 
 from edugrad.dtypes import dtypes
@@ -36,6 +36,16 @@
 def __getitem__(
     tensor: "Tensor", val
 ) -> "Tensor":  # val: Union[int, slice, Tensor, None, Ellipsis, Tuple[Union[int, slice, Tensor, None, Ellipsis], ...]]
+    """Retrieves an element or a slice from the tensor based on the specified value.
+
+    Args:
+    - tensor (Tensor): The tensor from which to retrieve the element or slice.
+    - val: The index or slice object. Can be an integer, slice, Tensor, None, Ellipsis, or a combination thereof in a tuple.
+
+    Returns:
+    - Tensor: A tensor containing the retrieved element or slice.
+
+    """
     from edugrad.tensor import Tensor
 
     def normalize_int(e, i, dim_sz):
@@ -138,11 +148,27 @@ def normalize_int(e, i, dim_sz):
 
 
 def __setitem__(tensor: "Tensor", s, v):
+    """Assigns a value to a specified element or slice of the tensor.
+
+    Args:
+    - tensor (Tensor): The tensor to modify.
+    - s: The index or slice where the value will be assigned.
+    - v: The value to be assigned.
+
+    """
     return tensor.__getitem__(s).assign(v)
 
 
 # NOTE: using slice is discouraged and things should migrate to pad and shrink
 def tslice(tensor: "Tensor", arg: Sequence[Optional[Tuple[int, shape_int]]], value: float = 0) -> "Tensor":
+    """Applies slicing to a tensor, using padding and shrinking for manipulation.
+
+    Args:
+    - tensor (Tensor): The tensor to slice.
+    - arg (Sequence[Optional[Tuple[int, shape_int]]]): A sequence of tuples defining the slicing parameters.
+    - value (float): The padding value to be used if necessary.
+
+    """
     from edugrad.tensor import Tensor
 
     arg_ = tuple([a if a is not None else (0, s) for s, a in zip(tensor.shape, arg)])
@@ -154,7 +180,15 @@ def tslice(tensor: "Tensor", arg: Sequence[Optional[Tuple[int, shape_int]]], val
 
 
 def gather(tensor: "Tensor", idx: "Tensor", dim: int) -> "Tensor":
-    from edugrad._tensor import Tensor
+    """Gathers elements from the tensor along a specified dimension, according to indices specified in another tensor.
+
+    Args:
+    - tensor (Tensor): The tensor from which to gather elements.
+    - idx (Tensor): The tensor containing indices to gather.
+    - dim (int): The dimension along which to gather.
+
+    """
+    from edugrad.tensor import Tensor
 
     assert idx.ndim == tensor.ndim, "tensor.ndim must equal idx.ndim"
     assert all(