Skip to content

TensorShare

Bases: BaseModel

Base model for tensor sharing.

Source code in src/tensorshare/schema.py
class TensorShare(BaseModel):
    """Base model for tensor sharing."""

    tensors: bytes
    size: ByteSize

    model_config = ConfigDict(
        extra="forbid",
        frozen=True,
        json_schema_extra={
            "tensors": b"SAAAAAAAAAB7ImVtYmVkZGluZ3MiOnsiZHR5cGUiOiJGMzIiLCJzaGFwZSI6WzEsMV0sImRhdGFfb2Zmc2V0cyI6WzAsNF19fSAgICAgICAAAAAA",
            "size": 84,
        },
    )

    @property
    def size_as_str(self) -> str:
        """Return size as a string in human readable format."""
        return self.size.human_readable()

    @classmethod
    def from_dict(
        cls,
        tensors: Dict[
            str,
            Union["Array", "np.ndarray", "paddle.Tensor", "tf.Tensor", "torch.Tensor"],
        ],
        metadata: Optional[Dict[str, str]] = None,
        backend: Optional[Union[str, Backend]] = None,
    ) -> TensorShare:
        """
        Create a TensorShare object from a dictionary of tensors.

        Args:
            tensors (Dict[str, Union[Array, np.ndarray, paddle.Tensor, tf.Tensor, torch.Tensor]]):
                Tensors stored in a dictionary with their name as key.
            metadata (Optional[Dict[str, str]], optional):
                Metadata to add to the safetensors file. Defaults to None.
            backend (Optional[Union[str, Backend]], optional):
                Backend to use for the conversion. Defaults to None.

        Raises:
            TypeError: If tensors is not a dictionary.
            TypeError: If backend is not a string or an instance of Backend enum.
            ValueError: If backend is not supported.
            ValueError: If tensors is empty.

        Returns:
            TensorShare: A TensorShare object.
        """
        _tensors, size = TensorProcessor.serialize(
            tensors=tensors, metadata=metadata, backend=backend
        )

        return cls(tensors=_tensors, size=size)

    def to_tensors(
        self,
        backend: Union[str, Backend],
    ) -> Dict[
        str, Union["Array", "np.ndarray", "paddle.Tensor", "tf.Tensor", "torch.Tensor"]
    ]:
        """
        Convert a TensorShare object to a dictionary of tensors in the specified backend.

        Args:
            backend (Union[str, Backend]): Backend to use for the conversion.

        Raises:
            TypeError: If backend is not a string or an instance of Backend enum.
            ValueError: If backend is not supported.

        Returns:
            Dict[str, Union[Array, np.ndarray, paddle.Tensor, tf.Tensor, torch.Tensor]]:
                Tensors stored in a dictionary with their name as key in the specified backend.
        """
        return TensorProcessor.deserialize(self.tensors, backend=backend)

size_as_str: str property

Return size as a string in human readable format.

from_dict(tensors, metadata=None, backend=None) classmethod

Create a TensorShare object from a dictionary of tensors.

Parameters:

Name Type Description Default
tensors Dict[str, Union[Array, ndarray, Tensor, Tensor, Tensor]]

Tensors stored in a dictionary with their name as key.

required
metadata Optional[Dict[str, str]]

Metadata to add to the safetensors file. Defaults to None.

None
backend Optional[Union[str, Backend]]

Backend to use for the conversion. Defaults to None.

None

Raises:

Type Description
TypeError

If tensors is not a dictionary.

TypeError

If backend is not a string or an instance of Backend enum.

ValueError

If backend is not supported.

ValueError

If tensors is empty.

Returns:

Name Type Description
TensorShare TensorShare

A TensorShare object.

Source code in src/tensorshare/schema.py
@classmethod
def from_dict(
    cls,
    tensors: Dict[
        str,
        Union["Array", "np.ndarray", "paddle.Tensor", "tf.Tensor", "torch.Tensor"],
    ],
    metadata: Optional[Dict[str, str]] = None,
    backend: Optional[Union[str, Backend]] = None,
) -> TensorShare:
    """
    Create a TensorShare object from a dictionary of tensors.

    Args:
        tensors (Dict[str, Union[Array, np.ndarray, paddle.Tensor, tf.Tensor, torch.Tensor]]):
            Tensors stored in a dictionary with their name as key.
        metadata (Optional[Dict[str, str]], optional):
            Metadata to add to the safetensors file. Defaults to None.
        backend (Optional[Union[str, Backend]], optional):
            Backend to use for the conversion. Defaults to None.

    Raises:
        TypeError: If tensors is not a dictionary.
        TypeError: If backend is not a string or an instance of Backend enum.
        ValueError: If backend is not supported.
        ValueError: If tensors is empty.

    Returns:
        TensorShare: A TensorShare object.
    """
    _tensors, size = TensorProcessor.serialize(
        tensors=tensors, metadata=metadata, backend=backend
    )

    return cls(tensors=_tensors, size=size)

to_tensors(backend)

Convert a TensorShare object to a dictionary of tensors in the specified backend.

Parameters:

Name Type Description Default
backend Union[str, Backend]

Backend to use for the conversion.

required

Raises:

Type Description
TypeError

If backend is not a string or an instance of Backend enum.

ValueError

If backend is not supported.

Returns:

Type Description
Dict[str, Union['Array', 'np.ndarray', 'paddle.Tensor', 'tf.Tensor', 'torch.Tensor']]

Dict[str, Union[Array, np.ndarray, paddle.Tensor, tf.Tensor, torch.Tensor]]: Tensors stored in a dictionary with their name as key in the specified backend.

Source code in src/tensorshare/schema.py
def to_tensors(
    self,
    backend: Union[str, Backend],
) -> Dict[
    str, Union["Array", "np.ndarray", "paddle.Tensor", "tf.Tensor", "torch.Tensor"]
]:
    """
    Convert a TensorShare object to a dictionary of tensors in the specified backend.

    Args:
        backend (Union[str, Backend]): Backend to use for the conversion.

    Raises:
        TypeError: If backend is not a string or an instance of Backend enum.
        ValueError: If backend is not supported.

    Returns:
        Dict[str, Union[Array, np.ndarray, paddle.Tensor, tf.Tensor, torch.Tensor]]:
            Tensors stored in a dictionary with their name as key in the specified backend.
    """
    return TensorProcessor.deserialize(self.tensors, backend=backend)

Last update: 2023-08-20
Created: 2023-08-20