def rotate_weight(
    weight: Union[torch.Tensor, nn.Parameter],
    Q: Union[torch.Tensor, nn.Parameter],
    pre: bool = True,
    device: Optional[torch.device] = None,
) -> Union[torch.Tensor, nn.Parameter]:
    """Rotate the weight matrix by Q.

    :param weight: weight matrix. Shape is (out_features, in_features). Equivalent to W^T in y = xW + b
    :param Q: rotation matrix. If the Q dimension is different from the weight dimension, head-wise rotation
        is performed.
    :pre: whether to apply the rotation before the linear operation.
        True for pre-rotation: y' = (xQ^-1)W + b = x(Q^TW) + b
        False for post-rotation: y' = yQ = x(WQ) + bQ
    :param device: device to use for the computation.
    """
    dtype = weight.dtype
    original_device = weight.device
    out_features, in_features = weight.shape
    q_features = Q.shape[0]

    head_wise = (in_features != q_features) if pre else (out_features != q_features)

    # Convert to double precision, optionally move to device
    to_kwargs = {"dtype": torch.float64}
    if device is not None:
        to_kwargs["device"] = device
    weight = weight.to(**to_kwargs)
    Q = Q.to(**to_kwargs)

    if pre:
        # Q^T @ W = (W^T @ Q)^T = (weight @ Q)^T
        if head_wise:
            weight = weight.reshape(out_features, -1, q_features)
            weight = torch.matmul(weight, Q).reshape(out_features, -1)
        else:
            weight = torch.matmul(weight, Q)
    else:
        # W @ Q = (Q^T @ W^T)^T = (Q^T @ weight)^T
        if head_wise:
            weight = weight.t().reshape(in_features, -1, q_features)
            weight = torch.matmul(weight, Q).reshape(in_features, -1).t()
        else:
            weight = torch.matmul(Q.t(), weight)

    # Convert back to original precision and device
    to_kwargs = {"dtype": dtype}
    if device is not None:
        to_kwargs["device"] = original_device
    return weight.to(**to_kwargs)
Rotate the weight matrix by Q.

    :param weight: weight matrix. Shape is (out_features, in_features). Equivalent to W^T in y = xW + b
    :param Q: rotation matrix. If the Q dimension is different from the weight dimension, head-wise rotation
        is performed.
    :pre: whether to apply the rotation before the linear operation.
        True for pre-rotation: y' = (xQ^-1)W + b = x(Q^TW) + b
        False for post-rotation: y' = yQ = x(WQ) + bQ
    :param device: device to use for the computation.
    
olive/passes/pytorch/rotate.py