vortex_torch.utils

Classes

ElementwiseBinaryOpType(*values)
ElementwiseOpType(*values)
Mode(*values)
ReduceType(*values)

class Mode(*values)

Bases: Enum

execute = 1

class Schedule(*values)

Bases: Enum

W = 0

S = 1

class ReduceType(*values)

Bases: Enum

Mean = 0

Max = 1

Min = 2

L2Norm = 3

Sum = 4

class ElementwiseBinaryOpType(*values)

Bases: Enum

Maximum = 0

Minimum = 1

Add = 2

Mul = 3

WhereEqual = 4

WhereNotEqual = 5

WhereGreater = 6

WhereGreaterEqual = 7

WhereLess = 8

WhereLessEqual = 9

class ElementwiseOpType(*values)

Bases: Enum

Relu = 0

Sigmoid = 1

Silu = 2

Abs = 3

Add_Mul = 4

Log = 5

Exp = 6

class QuantizationType(*values)

Bases: Enum

BF16 = 0

FP8_E5M2 = 1

FP8_E4M3 = 2

indent_block(text, level=1)

Indent a multi-line text block by the given indentation level.

Parameters:

• text (str)

• level (int)

Return type:

str

DTYPE_STR_TO_TORCH = {'bfloat16': torch.bfloat16, 'float16': torch.float16, 'float32': torch.float32, 'fp8_e4m3fn': torch.float8_e4m3fn, 'fp8_e5m2': torch.float8_e5m2}

Canonical string -> torch.dtype. Keep this one-to-one; aliases go in _DTYPE_ALIASES so tooling that prints the valid set (choices=, error messages) can show canonical names without alias noise.

resolve_dtype(value, default=torch.bfloat16)

Resolve a dtype given either a torch.dtype or a canonical string.

Accepts None (returns default), a torch.dtype (returned as-is), a canonical string in DTYPE_STR_TO_TORCH, or a short alias in _DTYPE_ALIASES.

Parameters:

default (torch.dtype)

Return type:

torch.dtype

is_hopper_or_newer()

Check if the current CUDA device is Hopper architecture or newer.

is_hopper()

Check if the current CUDA device is exactly Hopper architecture.