def rearrange_tensor(tensor):
    """
    Rearranges the last dimension (D) of the input tensor based on the specified mapping:
    2d -> d, 2d+1 -> D/2 + d.

    Args:
        tensor (torch.Tensor): Input tensor of shape [B, H, L, D], where D is even.

    Returns:
        torch.Tensor: Tensor with rearranged last dimension, same shape as input.
    """
    B, H, L, D = tensor.shape
    if D % 2 != 0:
        raise ValueError("The last dimension D must be even.")

    half_D = D // 2
    indices = torch.empty(D, dtype=torch.long, device=tensor.device)

    # Fill the indices based on the mapping rule
    indices[:half_D] = torch.arange(0, D, 2, device=tensor.device)
    indices[half_D:] = torch.arange(1, D, 2, device=tensor.device)

    # Rearrange the tensor based on the computed indices
    return tensor.index_select(dim=-1, index=indices)

    Rearranges the last dimension (D) of the input tensor based on the specified mapping:
    2d -> d, 2d+1 -> D/2 + d.

    Args:
        tensor (torch.Tensor): Input tensor of shape [B, H, L, D], where D is even.

    Returns:
        torch.Tensor: Tensor with rearranged last dimension, same shape as input.
    
opensora/models/mmdit/math.py