import torch
from einops import rearrange


def unbind_time(x: torch.Tensor, d: int) -> torch.Tensor:
    return rearrange(x, "bs (t d) ... -> bs t d ...", d=d)


def bind_time(x: torch.Tensor) -> torch.Tensor:
    return rearrange(x, "bs t d ... -> bs (t d) ...")
