TinyTorch/tinytorch/core/tensor.py

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../modules/tensor/tensor_dev.ipynb.

# %% auto 0
__all__ = ['Tensor']

# %% ../../modules/tensor/tensor_dev.ipynb 1
import numpy as np
import sys
from typing import Union, List, Tuple, Optional, Any

# %% ../../modules/tensor/tensor_dev.ipynb 3
class Tensor:
    """
    TinyTorch Tensor: N-dimensional array with ML operations.
    
    The fundamental data structure for all TinyTorch operations.
    Wraps NumPy arrays with ML-specific functionality.
    """
    
    def __init__(self, data: Union[int, float, List, np.ndarray], dtype: Optional[str] = None):
        """
        Create a new tensor from data.
        
        Args:
            data: Input data (scalar, list, or numpy array)
            dtype: Data type ('float32', 'int32', etc.). Defaults to auto-detect.
        """
        # Convert input to numpy array
        if isinstance(data, (int, float, np.number)):
            # Handle Python and NumPy scalars
            if dtype is None:
                # Auto-detect type: int for integers, float32 for floats
                if isinstance(data, int) or (isinstance(data, np.number) and np.issubdtype(type(data), np.integer)):
                    dtype = 'int32'
                else:
                    dtype = 'float32'
            self._data = np.array(data, dtype=dtype)
        elif isinstance(data, list):
            # Let NumPy auto-detect type, then convert if needed
            temp_array = np.array(data)
            if dtype is None:
                # Keep NumPy's auto-detected type, but prefer common ML types
                if np.issubdtype(temp_array.dtype, np.integer):
                    dtype = 'int32'
                elif np.issubdtype(temp_array.dtype, np.floating):
                    dtype = 'float32'
                else:
                    dtype = temp_array.dtype
            self._data = temp_array.astype(dtype)
        elif isinstance(data, np.ndarray):
            self._data = data.astype(dtype or data.dtype)
        else:
            raise TypeError(f"Cannot create tensor from {type(data)}")
    
    @property
    def data(self) -> np.ndarray:
        """Access underlying numpy array."""
        return self._data
    
    @property
    def shape(self) -> Tuple[int, ...]:
        """Get tensor shape."""
        return self._data.shape
    
    @property
    def size(self) -> int:
        """Get total number of elements."""
        return self._data.size
    
    @property
    def dtype(self) -> np.dtype:
        """Get data type as numpy dtype."""
        return self._data.dtype
    
    def __repr__(self) -> str:
        """String representation."""
        return f"Tensor({self._data.tolist()}, shape={self.shape}, dtype={self.dtype})"

# %% ../../modules/tensor/tensor_dev.ipynb 6
# Add arithmetic operations to the Tensor class
def _add_arithmetic_ops():
    """Add arithmetic operations to Tensor class."""
    
    def __add__(self, other: Union['Tensor', int, float]) -> 'Tensor':
        """Addition: tensor + other"""
        if isinstance(other, Tensor):
            return Tensor(self._data + other._data)
        else:  # scalar
            return Tensor(self._data + other)
    
    def __sub__(self, other: Union['Tensor', int, float]) -> 'Tensor':
        """Subtraction: tensor - other"""
        if isinstance(other, Tensor):
            return Tensor(self._data - other._data)
        else:  # scalar
            return Tensor(self._data - other)
    
    def __mul__(self, other: Union['Tensor', int, float]) -> 'Tensor':
        """Multiplication: tensor * other"""
        if isinstance(other, Tensor):
            return Tensor(self._data * other._data)
        else:  # scalar
            return Tensor(self._data * other)
    
    def __truediv__(self, other: Union['Tensor', int, float]) -> 'Tensor':
        """Division: tensor / other"""
        if isinstance(other, Tensor):
            return Tensor(self._data / other._data)
        else:  # scalar
            return Tensor(self._data / other)
    
    def __radd__(self, other: Union[int, float]) -> 'Tensor':
        """Reverse addition: scalar + tensor"""
        return Tensor(other + self._data)
    
    def __rmul__(self, other: Union[int, float]) -> 'Tensor':
        """Reverse multiplication: scalar * tensor"""
        return Tensor(other * self._data)
    
    # Add methods to Tensor class
    Tensor.__add__ = __add__
    Tensor.__sub__ = __sub__
    Tensor.__mul__ = __mul__
    Tensor.__truediv__ = __truediv__
    Tensor.__radd__ = __radd__
    Tensor.__rmul__ = __rmul__

# Apply the arithmetic operations
_add_arithmetic_ops()

# %% ../../modules/tensor/tensor_dev.ipynb 9
# Add utility methods to the Tensor class
def _add_utility_methods():
    """Add utility methods to Tensor class."""
    
    def reshape(self, *shape: int) -> 'Tensor':
        """Reshape tensor to new dimensions."""
        return Tensor(self._data.reshape(shape))
    
    def transpose(self) -> 'Tensor':
        """Transpose the tensor (swap dimensions)."""
        return Tensor(self._data.T)
    
    def sum(self, axis: Optional[int] = None) -> 'Tensor':
        """Sum elements along axis (or all elements if axis=None)."""
        result = self._data.sum(axis=axis)
        return Tensor(result)
    
    def mean(self, axis: Optional[int] = None) -> 'Tensor':
        """Mean of elements along axis (or all elements if axis=None)."""
        result = self._data.mean(axis=axis)
        return Tensor(result)
    
    def max(self, axis: Optional[int] = None) -> 'Tensor':
        """Maximum element along axis (or all elements if axis=None)."""
        result = self._data.max(axis=axis)
        return Tensor(result)
    
    def min(self, axis: Optional[int] = None) -> 'Tensor':
        """Minimum element along axis (or all elements if axis=None)."""
        result = self._data.min(axis=axis)
        return Tensor(result)
    
    def item(self) -> Union[int, float]:
        """Convert single-element tensor to Python scalar."""
        if self.size != 1:
            raise ValueError(f"Cannot convert tensor of size {self.size} to scalar")
        return self._data.item()
    
    def numpy(self) -> np.ndarray:
        """Convert to numpy array."""
        return self._data.copy()
    
    # Add methods to Tensor class
    Tensor.reshape = reshape
    Tensor.transpose = transpose
    Tensor.sum = sum
    Tensor.mean = mean
    Tensor.max = max
    Tensor.min = min
    Tensor.item = item
    Tensor.numpy = numpy

# Apply the utility methods
_add_utility_methods()