tensor 🔗

def as_bytelike(self) -> bytes | bytearray | memoryview:
    """Returns the tensor data buffer as a byte-like object.

    Returns:
        The data of the tensor as a byte-like object.
    """
    if self._c_data is not None:
        if isinstance(self._c_data, CBytes):
            return self._c_data.value
        return self._c_data.as_memoryview()
    return CBytes.from_c_obj(
        self._c_ptr_or_raise.data, self._c_ptr_or_raise.size
    ).value

def as_memoryview(self) -> memoryview:
    """Returns the tensor data buffer as memoryview.

    Returns:
        The data of the tensor as memoryview.
    """
    data = self.as_bytelike()
    return data if isinstance(data, memoryview) else memoryview(data)

def as_numpy(self) -> "np.ndarray":
    """Returns the tensor data buffer as a NumPy array.

    Returns:
        The data of the tensor as a NumPy array.
    """
    if not HAS_NUMPY:
        msg = "NumPy is not available"
        raise NotImplementedError(msg)

    if isinstance(self._c_data, CNumpyArray):
        return self._c_data.value

    dtype = TensorTypeMapper.type_to_numpy(self.data_type)
    return CNumpyArray.from_c_obj(
        self._c_ptr_or_raise.data,
        self._c_ptr_or_raise.size,
        self.shape,
        dtype,
    ).value

def as_torch(self) -> "torch.Tensor":
    """Returns the tensor as a PyTorch tensor.

    Returns:
        The tensor as a PyTorch tensor.

    Raises:
        NotImplementedError: If PyTorch is not installed.
    """
    if not HAS_TORCH:
        msg = "PyTorch is not available"
        raise NotImplementedError(msg)

    if isinstance(self._data, torch.Tensor):
        return self._data

    dtype = TensorTypeMapper.type_to_torch(self.data_type)
    buf = ffi.buffer(self._c_ptr_or_raise.data, self._c_ptr_or_raise.size)

    return torch.frombuffer(buf, dtype=dtype).reshape(self.shape)

@classmethod
def empty(cls) -> "Tensor":
    """Initializes a new empty `Tensor` object.

    The object has a NULL pointer in place of the C struct.

    Returns:
        A new `Tensor` object
    """
    inst = cls.__new__(cls)
    inst._dims = None
    inst._data = None
    inst._data_type = None
    inst._c_data = None
    inst._c_obj_ptr = ffi.new("struct vaccel_torch_tensor **")
    inst._c_obj = inst._c_obj_ptr[0]
    inst._c_obj_ptr[0] = ffi.NULL
    inst._c_size = ffi.sizeof("struct vaccel_torch_tensor *")
    return inst

@classmethod
def from_buffer(
    cls,
    dims: list[int],
    data_type: TensorType,
    data: bytes | bytearray | memoryview,
) -> "Tensor":
    """Initializes a new `Tensor` object from byte-like data.

    Args:
        dims: The dims to be passed to the C struct.
        data_type: The data_type to be passed to the C struct.
        data: The data to be passed to the C struct.

    Returns:
        A new `Tensor` object
    """
    inst = cls.__new__(cls)
    inst._dims = dims
    inst._data = data
    inst._data_type = data_type
    inst._c_data = CBytes(inst._data)
    inst._c_obj_ptr = ffi.NULL
    inst._c_obj_data = inst._c_data._c_ptr
    super().__init__(inst)
    return inst

@classmethod
def from_c_obj(cls, c_obj: ffi.CData) -> "Tensor":
    """Initializes a new `Tensor` object from an existing C struct.

    Args:
        c_obj: A pointer to a `struct vaccel_torch_tensor` C object.

    Returns:
        A new `Tensor` object
    """
    type_str = ffi.getctype(ffi.typeof(c_obj))
    if type_str != "struct vaccel_torch_tensor *":
        msg = f"Expected 'struct vaccel_torch_tensor *', got '{type_str}'"
        raise TypeError(msg)

    inst = cls.__new__(cls)
    inst._dims = None
    inst._data = None
    inst._data_type = None
    inst._c_data = None
    inst._c_obj_ptr = ffi.NULL
    inst._c_obj_data = ffi.NULL
    inst._c_obj = c_obj
    inst._c_size = ffi.sizeof(inst._c_obj)
    return inst

@classmethod
def from_numpy(cls, data: "np.ndarray") -> "Tensor":
    """Initializes a new `Tensor` object from a NumPy array.

    Args:
        data: The NumPy array containing the tensor data.

    Returns:
        A new `Tensor` object

    Raises:
        NotImplementedError: If NumPy is not installed.
    """
    if not HAS_NUMPY:
        msg = "NumPy is not available"
        raise NotImplementedError(msg)

    inst = cls.__new__(cls)
    inst._dims = list(data.shape)
    inst._data = data
    inst._data_type = TensorTypeMapper.type_from_numpy(inst._data.dtype)
    inst._c_data = CNumpyArray(inst._data)
    inst._c_obj_ptr = ffi.NULL
    inst._c_obj_data = inst._c_data._c_ptr
    super().__init__(inst)
    return inst

@classmethod
def from_torch(cls, data: "torch.Tensor") -> "Tensor":
    """Initializes a new `Tensor` object from a PyTorch tensor.

    Args:
        data: The PyTorch tensor containing the tensor data.

    Returns:
        A new `Tensor` object

    Raises:
        NotImplementedError: If PyTorch is not installed.
    """
    if not HAS_TORCH:
        msg = "PyTorch is not available"
        raise NotImplementedError(msg)

    inst = cls.__new__(cls)
    inst._dims = list(data.shape)
    inst._data = data.contiguous().cpu()
    inst._data_type = TensorTypeMapper.type_from_torch(inst._data.dtype)
    inst._c_obj_ptr = ffi.NULL
    data_ptr_int = inst._data.untyped_storage().data_ptr()
    data_ptr = ffi.cast("void *", data_ptr_int)
    data_size = inst._data.untyped_storage().nbytes()
    inst._c_data = CBytes.from_c_obj(data_ptr, data_size)
    inst._c_obj_data = inst._c_data._c_ptr
    super().__init__(inst)
    return inst

def to_bytes(self) -> bytes:
    """Returns the tensor data buffer as bytes.

    Returns:
        The data of the tensor as bytes.
    """
    data = self.as_bytelike()
    return data if isinstance(data, bytes) else bytes(data)

@classmethod
def type_from_numpy(cls, dtype: "np.dtype") -> TensorType:
    """Converts a NumPy `dtype` to `TensorType`.

    Args:
        dtype: A NumPy `dtype` object or something convertible to
            `np.dtype`.

    Returns:
        A corresponding tensor type value.

    Raises:
        NotImplementedError: If NumPy is not installed.
        ValueError: If the `dtype` value is not supported.
    """
    if not HAS_NUMPY:
        msg = "NumPy is not available"
        raise NotImplementedError(msg)

    dtype = np.dtype(dtype)
    if dtype not in cls._NUMPY_TO_TENSOR_TYPE:
        supported = ", ".join(str(d) for d in cls._NUMPY_TO_TENSOR_TYPE)
        msg = f"Unsupported NumPy dtype: {dtype}. Supported: {supported}"
        raise ValueError(msg)

    return cls._NUMPY_TO_TENSOR_TYPE[dtype]

@classmethod
def type_from_torch(cls, dtype: "torch.dtype") -> TensorType:
    """Converts a PyTorch `dtype` to `TensorType`.

    Args:
        dtype: A PyTorch `dtype` object.

    Returns:
        A corresponding tensor type value.

    Raises:
        NotImplementedError: If PyTorch is not installed.
        ValueError: If the `dtype` value is not supported.
    """
    if not HAS_TORCH:
        msg = "PyTorch is not available"
        raise NotImplementedError(msg)

    if dtype not in cls._TORCH_TO_TENSOR_TYPE:
        supported = ", ".join(str(d) for d in cls._TORCH_TO_TENSOR_TYPE)
        msg = f"Unsupported PyTorch dtype: {dtype}. Supported: {supported}"
        raise ValueError(msg)

    return cls._TORCH_TO_TENSOR_TYPE[dtype]

@classmethod
def type_to_c_size(cls, tensor_type: TensorType) -> int:
    """Converts a `TensorType` to a C type size (in bytes).

    Args:
        tensor_type: The tensor type value.

    Returns:
        A corresponding C type size in bytes.

    Raises:
        ValueError: If the `tensor_type` value is not supported.
    """
    if tensor_type not in cls._TENSOR_TYPE_TO_C:
        supported = ", ".join(str(d) for d in cls._TENSOR_TYPE_TO_C)
        msg = (
            f"Unsupported TensorType: {tensor_type}. Supported: {supported}"
        )
        raise ValueError(msg)
    return cls._TENSOR_TYPE_TO_C[tensor_type][1]

@classmethod
def type_to_c_type(cls, tensor_type: TensorType) -> str:
    """Converts a `TensorType` to a C type string.

    Args:
        tensor_type: The tensor type value.

    Returns:
        A corresponding C type as a string (e.g., "float", "int64_t").

    Raises:
        ValueError: If the `tensor_type` value is not supported.
    """
    if tensor_type not in cls._TENSOR_TYPE_TO_C:
        supported = ", ".join(str(d) for d in cls._TENSOR_TYPE_TO_C)
        msg = (
            f"Unsupported TensorType: {tensor_type}. Supported: {supported}"
        )
        raise ValueError(msg)
    return cls._TENSOR_TYPE_TO_C[tensor_type][0]

@classmethod
def type_to_numpy(cls, ttype: TensorType) -> "np.dtype":
    """Converts a `TensorType` to a NumPy `dtype`.

    Args:
        ttype: A `TensorType` enum value.

    Returns:
        A corresponding NumPy `dtype` object.

    Raises:
        NotImplementedError: If NumPy is not installed.
        ValueError: If the `TensorType` value is not supported.
    """
    if not HAS_NUMPY:
        msg = "NumPy is not available"
        raise NotImplementedError(msg)

    if ttype not in cls._TENSOR_TYPE_TO_NUMPY:
        supported = ", ".join(str(t) for t in cls._TENSOR_TYPE_TO_NUMPY)
        msg = f"Unsupported TensorType: {ttype}. Supported: {supported}"
        raise ValueError(msg)

    return cls._TENSOR_TYPE_TO_NUMPY[ttype]

@classmethod
def type_to_torch(cls, ttype: TensorType) -> "torch.dtype":
    """Converts a `TensorType` to a PyTorch `dtype`.

    Args:
        ttype: A `TensorType` enum value.

    Returns:
        A corresponding PyTorch `dtype` object.

    Raises:
        NotImplementedError: If PyTorch is not installed.
        ValueError: If the `TensorType` value is not supported.
    """
    if not HAS_TORCH:
        msg = "PyTorch is not available"
        raise NotImplementedError(msg)

    if ttype not in cls._TENSOR_TYPE_TO_TORCH:
        supported = ", ".join(str(t) for t in cls._TENSOR_TYPE_TO_TORCH)
        msg = f"Unsupported TensorType: {ttype}. Supported: {supported}"
        raise ValueError(msg)

    return cls._TENSOR_TYPE_TO_TORCH[ttype]