quark.torch.algorithm.rotation.rotation_utils

Module Contents

Classes

RMSNorm

Root Mean Square Layer Normalization (RMSNorm).

Functions

rotate_in_channels(→ None)	Rotate the input channels of a weight matrix.
rotate_out_channels(→ None)	Rotate the output channels of a weight matrix.
get_rotation_matrix(→ torch.Tensor)	Get a random rotation matrix for the given number of channels.

class quark.torch.algorithm.rotation.rotation_utils.RMSNorm(hidden_size: int, eps: float = 1e-06)

Root Mean Square Layer Normalization (RMSNorm).

forward(hidden_states: torch.Tensor) → torch.Tensor

Apply RMSNorm normalization to hidden states.

quark.torch.algorithm.rotation.rotation_utils.rotate_in_channels(module: torch.nn.Module, rotation: torch.Tensor) → None

Rotate the input channels of a weight matrix.

quark.torch.algorithm.rotation.rotation_utils.rotate_out_channels(module: torch.nn.Module, rotation: torch.Tensor) → None

Rotate the output channels of a weight matrix.

quark.torch.algorithm.rotation.rotation_utils.get_rotation_matrix(num_channels: int, random: bool = True) → torch.Tensor

Get a random rotation matrix for the given number of channels.