For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

get_safetensors_idx

def get_safetensors_idx(head_dim_idx: Int, head_size: Int) -> Tuple[Int, Int]

Returns the real and imaginary element indices for a safetensors layout.

Safetensors stores RoPE weights as all real components followed by all imaginary components, so the imaginary index is offset by half the head size.

Args:

head_dim_idx (Int): The dimension index within the head.
head_size (Int): The total size of the head dimension.

Returns:

Tuple[Int, Int]: A tuple of the real index and the imaginary index.