import torch


def initialize_pos_array(n_frames, init, device, init_frame=0):
    """ Constructs empty array to be filled with positions. Init_frame sets which frame should be used for init. infers shapes from init.
    Return shsape (n_frames, init.shape[0], init.shape[1])."""
    pos_array = torch.zeros((n_frames, init.shape[0], init.shape[1]), device=device)
    pos_array[init_frame] = torch.tensor(init, dtype=torch.float32)
    return pos_array