synor/sdk/python/synor_compute/tensor.py

"""Tensor utilities for Synor Compute SDK."""

from typing import Any
import base64
import numpy as np

from .types import Precision


class Tensor:
    """
    Tensor wrapper for compute operations.

    Example:
        >>> tensor = Tensor.from_numpy(np.random.randn(1024, 1024))
        >>> tensor = Tensor.zeros((1024, 1024), dtype=Precision.FP16)
    """

    def __init__(
        self,
        data: np.ndarray,
        dtype: Precision = Precision.FP32,
    ):
        self._data = data
        self._dtype = dtype

    @classmethod
    def from_numpy(cls, array: np.ndarray, dtype: Precision | None = None) -> "Tensor":
        """Create a tensor from a numpy array."""
        if dtype is None:
            dtype = cls._infer_dtype(array.dtype)
        return cls(array, dtype)

    @classmethod
    def zeros(cls, shape: tuple[int, ...], dtype: Precision = Precision.FP32) -> "Tensor":
        """Create a tensor of zeros."""
        np_dtype = cls._to_numpy_dtype(dtype)
        return cls(np.zeros(shape, dtype=np_dtype), dtype)

    @classmethod
    def ones(cls, shape: tuple[int, ...], dtype: Precision = Precision.FP32) -> "Tensor":
        """Create a tensor of ones."""
        np_dtype = cls._to_numpy_dtype(dtype)
        return cls(np.ones(shape, dtype=np_dtype), dtype)

    @classmethod
    def random(cls, shape: tuple[int, ...], dtype: Precision = Precision.FP32) -> "Tensor":
        """Create a tensor with random values."""
        np_dtype = cls._to_numpy_dtype(dtype)
        return cls(np.random.random(shape).astype(np_dtype), dtype)

    @classmethod
    def randn(cls, shape: tuple[int, ...], dtype: Precision = Precision.FP32) -> "Tensor":
        """Create a tensor with normal distribution."""
        np_dtype = cls._to_numpy_dtype(dtype)
        return cls(np.random.randn(*shape).astype(np_dtype), dtype)

    @property
    def data(self) -> np.ndarray:
        """Get the underlying numpy array."""
        return self._data

    @property
    def shape(self) -> tuple[int, ...]:
        """Get tensor shape."""
        return self._data.shape

    @property
    def dtype(self) -> Precision:
        """Get tensor dtype."""
        return self._dtype

    @property
    def size(self) -> int:
        """Get total number of elements."""
        return self._data.size

    @property
    def ndim(self) -> int:
        """Get number of dimensions."""
        return self._data.ndim

    @property
    def nbytes(self) -> int:
        """Get byte size."""
        return self._data.nbytes

    def reshape(self, shape: tuple[int, ...]) -> "Tensor":
        """Reshape tensor."""
        return Tensor(self._data.reshape(shape), self._dtype)

    def to(self, dtype: Precision) -> "Tensor":
        """Convert to different precision."""
        if dtype == self._dtype:
            return self
        np_dtype = self._to_numpy_dtype(dtype)
        return Tensor(self._data.astype(np_dtype), dtype)

    def numpy(self) -> np.ndarray:
        """Convert to numpy array."""
        return self._data

    def serialize(self) -> dict[str, Any]:
        """Serialize tensor for transmission."""
        data_bytes = self._data.tobytes()
        data_b64 = base64.b64encode(data_bytes).decode("ascii")
        return {
            "data": data_b64,
            "shape": list(self._data.shape),
            "dtype": self._dtype.value,
        }

    @classmethod
    def deserialize(cls, data: dict[str, Any]) -> "Tensor":
        """Deserialize tensor from transmission format."""
        data_bytes = base64.b64decode(data["data"])
        dtype = Precision(data["dtype"])
        np_dtype = cls._to_numpy_dtype(dtype)
        array = np.frombuffer(data_bytes, dtype=np_dtype).reshape(data["shape"])
        return cls(array, dtype)

    @staticmethod
    def _infer_dtype(np_dtype: np.dtype) -> Precision:
        """Infer Precision from numpy dtype."""
        if np_dtype == np.float64:
            return Precision.FP64
        elif np_dtype == np.float32:
            return Precision.FP32
        elif np_dtype == np.float16:
            return Precision.FP16
        elif np_dtype == np.int8:
            return Precision.INT8
        else:
            return Precision.FP32

    @staticmethod
    def _to_numpy_dtype(dtype: Precision) -> np.dtype:
        """Convert Precision to numpy dtype."""
        mapping = {
            Precision.FP64: np.float64,
            Precision.FP32: np.float32,
            Precision.FP16: np.float16,
            Precision.BF16: np.float16,  # Approximate
            Precision.INT8: np.int8,
            Precision.INT4: np.int8,  # Approximate
        }
        return np.dtype(mapping.get(dtype, np.float32))

    def __repr__(self) -> str:
        return f"Tensor(shape={self.shape}, dtype={self._dtype.value})"

    def __str__(self) -> str:
        return f"Tensor({self._data}, dtype={self._dtype.value})"