Skip to content

torchmx.utils

get_logger

[view source] 

1
2
3
4
def get_logger(logger_name: str = "TORCHMX",
               format_string:
               str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
               console_output: bool = True) -> logging.Logger

Returns a logger with the specified name and format.

Arguments:

  • logger_name str, optional - Name of the logger. Defaults to "TORCHMX".
  • format_string str, optional - Format of the log message. Defaults to "%(asctime)s - %(name)s - %(levelname)s - %(message)s".
  • console_output bool, optional - Whether to output the logs to the console. Defaults to True.

Returns:

  • logging.Logger - Logger with the specified name and format.

get_uniform_random_number

[view source] 

1
2
def get_uniform_random_number(min_val: int, max_val: int, shape: Iterable[int],
                              dtype: torch.dtype) -> torch.Tensor

Generate random numbers from uniform distribution of range [min_val, max_val)

Arguments:

  • min_val int - minimum value of the range
  • max_val int - maximum value of the range
  • shape Iterable[int] - shape of the tensor
  • dtype torch.dtype - data type of the tensor

Returns:

  • torch.Tensor - tensor of shape shape and dtype dtype with random numbers

tensor_size_hp_to_fp4x2

[view source] 

1
2
def tensor_size_hp_to_fp4x2(orig_size: torch.Size,
                            packing_dim: int) -> List[int]

Converts the size of a tensor from half precision to fp4x2 precision.

Arguments:

  • orig_size torch.Size - The size of the original tensor.
  • packing_dim int - The dimension where for packing 2xuint4 per byte

Returns:

  • List[int] - The size of the tensor in fp4x2 precision.

tensor_size_fp4x2_to_hp

[view source] 

1
2
def tensor_size_fp4x2_to_hp(orig_size: torch.Size,
                            unpacking_dim: int) -> List[int]

Converts the size of a tensor from fp4x2 precision to half precision by unpacking the 4-bits into 8-bits.

Arguments:

  • orig_size torch.Size - The size of the original tensor.
  • unpacking_dim int - The dimension where for unpacking the uint4 values to a single byte

Returns:

  • List[int] - The size of the tensor in half precision.

unpack_uint4

[view source] 

1
2
def unpack_uint4(uint8_data: torch.Tensor,
                 packing_dim: int = -1) -> torch.Tensor

Unpacks a tensor of uint8 values into two tensors of uint4 values.

Arguments:

  • uint8_data torch.Tensor - A tensor containing packed uint8 values.
  • packing_dim int - The dimension along which the unpacking is performed.

Returns:

  • torch.Tensor - A tensor containing the unpacked uint4 values.

pack_uint4

[view source] 

1
2
def pack_uint4(uint8_data: torch.Tensor,
               packing_dim: int = -1) -> torch.Tensor

Packs uint4 data to unit8 format along the specified dimension.

Arguments:

  • uint4_data torch.Tensor - The input tensor containing uint8 data.
  • packing_dim int - The dimension along which to pack the data.

Returns:

  • torch.Tensor - A tensor with the packed uint4 data.

Raises:

  • AssertionError - If the size of the specified dimension is not even.

Notes:

The function assumes that the input data is contiguous and reshapes it accordingly. The packing is done by combining pairs of uint8 values into a single uint8 value where each original uint8 value is treated as a uint4.

set_seed

[view source] 

1
def set_seed(seed: int) -> None

Set the random seed for reproducibility.

Arguments:

  • seed int - The seed value to set.

Returns:

None