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

@@ -70,7 +70,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.
-    
+
     Args:
         tensor (Tensor): The reference tensor to check compatibility.
         x (Tensor | float): The tensor or scalar to be converted.

edugrad/_tensor/tensor_create.py

@@ -1,7 +1,8 @@
 """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.
+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
@@ -19,9 +20,9 @@
 # creation low-level op entrypoint *****
 
 
-def _loadop(op: LoadOps, sz: int, dtype: Optional[DType] = None, arg: Any = None, **kwargs) -> Tensor:
+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.
@@ -31,6 +32,7 @@ def _loadop(op: LoadOps, sz: int, dtype: Optional[DType] = None, arg: Any = None
 
     Returns:
     - Tensor: A new tensor created with the specified operation.
+
     """
     from edugrad.tensor import Tensor
 
@@ -110,6 +112,7 @@ def arange(start: int | float, stop: int | float | None, step: int | float, **kw
 
     Returns:
     - Tensor: A 1D tensor containing a sequence of numbers.
+
     """
     from edugrad.tensor import Tensor
 
@@ -193,6 +196,7 @@ 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
 

edugrad/_tensor/tensor_index_slice.py

@@ -44,6 +44,7 @@ def __getitem__(
 
     Returns:
     - Tensor: A tensor containing the retrieved element or slice.
+
     """
     from edugrad.tensor import Tensor
 
@@ -166,6 +167,7 @@ def tslice(tensor: "Tensor", arg: Sequence[Optional[Tuple[int, shape_int]]], val
     - 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
 
@@ -184,6 +186,7 @@ def gather(tensor: "Tensor", idx: "Tensor", dim: int) -> "Tensor":
     - 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
 

edugrad/_tensor/tensor_reshape.py

@@ -7,7 +7,7 @@
 import edugrad.function as function
 
 
-def reshape(tensor: Tensor, shape: Union[int, Tuple[int, ...]], *args) -> Tensor:
+def reshape(tensor: Tensor, shape: int | tuple[int, ...], *args) -> Tensor:
     """Reshapes a tensor to the specified new shape.
 
     Args:
@@ -25,8 +25,9 @@ def reshape(tensor: Tensor, shape: Union[int, Tuple[int, ...]], *args) -> Tensor
     return function.Reshape.apply(tensor, shape=adjusted_shape)
 
 
-def expand(tensor: Tensor, shape: Union[int, Tuple[int, ...]], *args) -> Tensor:
-    """Expands the size of the tensor to the specified shape. -1 in the shape means the corresponding dimension is unchanged.
+def expand(tensor: Tensor, shape: int | tuple[int, ...], *args) -> Tensor:
+    """Expands the size of the tensor to the specified shape. -1 in the shape means the corresponding dimension is
+    unchanged.
 
     Args:
         tensor: The tensor to expand.
@@ -40,9 +41,8 @@ def expand(tensor: Tensor, shape: Union[int, Tuple[int, ...]], *args) -> Tensor:
     return function.Expand.apply(tensor, shape=expanded_shape)
 
 
-def permute(tensor: Tensor, order: Union[int, Tuple[int, ...]], *args) -> Tensor:
-    """
-    Permutes the tensor dimensions according to the specified order.
+def permute(tensor: Tensor, order: int | tuple[int, ...], *args) -> Tensor:
+    """Permutes the tensor dimensions according to the specified order.
 
     Args:
         tensor: The tensor to permute.
@@ -54,9 +54,8 @@ def permute(tensor: Tensor, order: Union[int, Tuple[int, ...]], *args) -> Tensor
     return function.Permute.apply(tensor, order=new_order)
 
 
-def flip(tensor: Tensor, axis: Union[int, List[int]], *args) -> Tensor:
-    """
-    Flips the tensor along the specified axes.
+def flip(tensor: Tensor, axis: int | list[int], *args) -> Tensor:
+    """Flips the tensor along the specified axes.
 
     Args:
         tensor: The tensor to flip.
@@ -69,9 +68,8 @@ def flip(tensor: Tensor, axis: Union[int, List[int]], *args) -> Tensor:
     return function.Flip.apply(tensor, axis=normalized_axes)
 
 
-def shrink(tensor: Tensor, arg: Tuple[Tuple[shape_int, shape_int] | None, ...]) -> Tensor:
-    """
-    Shrinks the tensor along each dimension according to the specified start and end indices.
+def shrink(tensor: Tensor, arg: tuple[tuple[shape_int, shape_int] | None, ...]) -> Tensor:
+    """Shrinks the tensor along each dimension according to the specified start and end indices.
 
     Args:
         tensor: The tensor to shrink.
@@ -86,9 +84,8 @@ def shrink(tensor: Tensor, arg: Tuple[Tuple[shape_int, shape_int] | None, ...])
     return tensor
 
 
-def pad(tensor: Tensor, arg: Tuple[Tuple[int, int] | None, ...], value: float) -> Tensor:
-    """
-    Pads the tensor along each dimension with the specified padding values.
+def pad(tensor: Tensor, arg: tuple[tuple[int, int] | None, ...], value: float) -> Tensor:
+    """Pads the tensor along each dimension with the specified padding values.
 
     Args:
         tensor: The tensor to pad.
@@ -97,6 +94,7 @@ def pad(tensor: Tensor, arg: Tuple[Tuple[int, int] | None, ...], value: float) -
 
     """
     from edugrad.tensor import Tensor
+
     # Determine padding for each dimension, defaulting to (0, 0)
     pad_arg = tuple(x if x is not None else (0, 0) for x in arg)
     # Apply padding operation if necessary
@@ -107,9 +105,8 @@ def pad(tensor: Tensor, arg: Tuple[Tuple[int, int] | None, ...], value: float) -
     return ret if value == 0 else ret + function.Pad.apply(Tensor.ones_like(tensor), arg=pad_arg).where(0, value)
 
 
-def pad2d(tensor: Tensor, padding: Union[List[int], Tuple[int, ...]], value: float) -> Tensor:
-    """
-    Pads the tensor with 2D padding specified for each side.
+def pad2d(tensor: Tensor, padding: list[int] | tuple[int, ...], value: float) -> Tensor:
+    """Pads the tensor with 2D padding specified for each side.
 
     Args:
         tensor: The tensor to pad.
@@ -123,8 +120,7 @@ def pad2d(tensor: Tensor, padding: Union[List[int], Tuple[int, ...]], value: flo
 
 
 def transpose(tensor: Tensor, ax1: int, ax2: int) -> Tensor:
-    """
-    Transposes two dimensions of the tensor.
+    """Transposes two dimensions of the tensor.
 
     Args:
         tensor: The tensor to transpose.
@@ -138,8 +134,7 @@ def transpose(tensor: Tensor, ax1: int, ax2: int) -> Tensor:
 
 
 def _flatten(tensor: Tensor, start_dim: int) -> Tensor:
-    """
-    Flattens the tensor from the specified start dimension.
+    """Flattens the tensor from the specified start dimension.
 
     Args:
         tensor: The tensor to flatten.
@@ -150,8 +145,7 @@ def _flatten(tensor: Tensor, start_dim: int) -> Tensor:
 
 
 def squeeze(tensor: Tensor, dim: Optional[int]) -> Tensor:
-    """
-    Squeezes the tensor by removing dimensions of size 1.
+    """Squeezes the tensor by removing dimensions of size 1.
 
     Args:
         tensor: The tensor to squeeze.
@@ -168,12 +162,15 @@ def squeeze(tensor: Tensor, dim: Optional[int]) -> Tensor:
         )
     if dim < 0:
         dim += tensor.ndim
-    return tensor if tensor.shape[dim] != 1 else tensor.reshape(*[size for idx, size in enumerate(tensor.shape) if idx != dim])
+    return (
+        tensor
+        if tensor.shape[dim] != 1
+        else tensor.reshape(*[size for idx, size in enumerate(tensor.shape) if idx != dim])
+    )
 
 
 def unsqueeze(tensor: Tensor, dim: int) -> Tensor:
-    """
-    Adds a dimension of size 1 to the tensor at the specified position.
+    """Adds a dimension of size 1 to the tensor at the specified position.
 
     Args:
         tensor: The tensor to unsqueeze.

edugrad/helpers.py

@@ -1,4 +1,4 @@
-"""Contains helper functions and DEBUG integer for verbose debugging used throughout the package."""
+"""Contains helper functions, a shape_int type and a DEBUG integer for verbose debugging used throughout the package."""
 
 from typing import Union, Tuple, Iterator, Any
 import os