def draw_video_boundary(x: jnp.ndarray) -> jnp.ndarray:
  """Pads the frame boundary with black pixels. x.shape=[bs, lt, lh, lw, 3]."""
  assert x.ndim == 5
  width = 1
  x = jnp.pad(
      x, [(0, 0), (0, 0), (0, 0), (width, width), (0, 0)],
      mode='constant',
      constant_values=0)
  x1 = einops.rearrange(x, 'device_bs t h w c -> device_bs h (t w) c')
  x1 = jnp.pad(
      x1, [(0, 0), (width, width), (0, 0), (0, 0)],
      mode='constant',
      constant_values=0)
  x2 = jnp.concatenate(x1, axis=0)
  return x2
Pads the frame boundary with black pixels. x.shape=[bs, lt, lh, lw, 3].
videogvt/train_lib/maskgvt_train_lib.py