module ptwm.utils

COMPRESSED_DTYPE = torch.uint8

COMPRESSION_METHOD = 'HUFFMAN'

def decode_ptwm_tensor(blob: bytes, name: str) -> torch.Tensor

def divide_int(tensor: torch.Tensor, divisor: float) -> torch.Tensor

class DType(Enum):

def __init__(self, torch_dtype: torch.dtype | None, numpy_dtype: type | None, python_dtype: type | None, dtype_str: str, code: int) -> None

BFLOAT16 = (torch.bfloat16, None, None, 'bfloat16', 6)

BOOL = (torch.bool, np.bool_, bool, 'bool', 24)

CDOUBLE = (torch.cdouble, np.complex128, complex, 'cdouble', 12)

CFLOAT = (torch.cfloat, np.complex64, complex, 'cfloat', 10)

CHALF = (torch.complex32, None, None, 'chalf', 8)

code = code

COMPLEX128 = (torch.complex128, np.complex128, complex, 'complex128', 11)

COMPLEX32 = (torch.complex32, None, None, 'complex32', 7)

COMPLEX64 = (torch.complex64, np.complex64, complex, 'complex64', 9)

dtype_str = dtype_str

FLOAT = (torch.float, np.float32, float, 'float', 2)

FLOAT16 = (torch.float16, np.float16, None, 'float16', 4)

FLOAT32 = (torch.float32, np.float32, float, 'float32', 1)

FLOAT4_E2M1FN_X2 = (getattr(torch, 'float4_e2m1fn_x2', None), None, None, 'float4_e2m1fn_x2', 31)

FLOAT64 = (torch.float64, np.float64, float, 'float64', 3)

FLOAT8_E4M3FN = (torch.float8_e4m3fn, None, None, 'float8_e4m3fn', 29)

FLOAT8_E5M2 = (torch.float8_e5m2, None, None, 'float8_e5m2', 30)

def from_code(cls, code: int) -> str

def from_dtype(cls, dtype: Any) -> DType

HALF = (torch.half, np.float16, None, 'half', 5)

INT = (torch.int32, np.int32, int, 'int', 21)

INT16 = (torch.int16, np.int16, None, 'int16', 18)

INT32 = (torch.int32, np.int32, int, 'int32', 20)

INT64 = (torch.int64, np.int64, int, 'int64', 22)

INT8 = (torch.int8, np.int8, None, 'int8', 17)

LONG = (torch.int64, np.int64, int, 'long', 23)

NONE = (None, None, None, 'none', 0)

numpy_dtype = numpy_dtype

python_dtype = python_dtype

QINT32 = (torch.qint32, None, None, 'qint32', 27)

QINT8 = (torch.qint8, None, None, 'qint8', 26)

QUINT4X2 = (torch.quint4x2, None, None, 'quint4x2', 28)

QUINT8 = (torch.quint8, None, None, 'quint8', 25)

SHORT = (torch.int16, np.int16, None, 'short', 19)

torch_dtype = torch_dtype

UINT16 = (None, np.uint16, None, 'uint16', 14)

UINT32 = (None, np.uint32, None, 'uint32', 15)

UINT64 = (None, np.uint64, None, 'uint64', 16)

UINT8 = (torch.uint8, np.uint8, None, 'uint8', 13)

def get_compressed_tensors_metadata(metadata: dict[str, str]) -> dict[str, CompressedTensorInfo]

def get_dtype_bits(dtype: torch.dtype) -> tuple[int, torch.dtype]

def is_floating_point(data_format_value: int, data: Any, bytearray_dtype: str | None) -> bool | None

def multi_process_patcher(patch_func: Callable[..., Any]) -> None

def multiply_if_max_below(tensor: torch.Tensor, max_val: float, multiplier: float, dtype: torch.dtype) -> tuple[torch.Tensor, bool]

def pack_shape(shape: torch.Size | list[int] | tuple[int, ...]) -> bytes

def raw_to_tensor(raw: bytes, shape: list[int] | tuple[int, ...] | None, dtype_name: str) -> torch.Tensor

def unpack_shape(packed_data: bytes) -> tuple[tuple[int, ...], int]