code: def extract_tubelets_from_video(
      x: jnp.ndarray,
      tubelet_size: Tuple[int, int, int],
      select_central_frame: bool,
  ) -> jnp.ndarray:
      """Extracts tubelets from videos for use as regression targets.
  
      Args:
        x: Input tensor of shape [batch, time, height, width, channels]
        tubelet_size: Tuple containing tubelet/patch size parameterised as
          [ph, pw, pt].
        select_central_frame: If True, select the central frame as the feature
          regression target.
  
      Returns:
        Tensor of shape [n, gt * gh * gw, pt * ph * pw * c] where
          gt = t // pt, gh = h // ph, gw = w // pw.
      """
      ph, pw, pt = tubelet_size
      n, t, h, w, c = x.shape
      gt, gh, gw = t // pt, h // ph, w // pw
      x = x.reshape([n, gt, pt, gh, ph, gw, pw, c])
      # Shape is then [n, gt, gh, gw, pt, ph, pw, c].
      x = jnp.transpose(x, axes=[0, 1, 3, 5, 2, 4, 6, 7])
      if select_central_frame:
          x = x[:, :, :, :, pt // 2, :, :, :]
          pt = 1
      return x.reshape([n, gt * gh * gw, pt * ph * pw * c])
  
file_path: scenic/scenic/projects/av_mae/base_model.py
imports: import functools
  from typing import Dict, Optional, Tuple, Union
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models import regression_model
  from scenic.model_lib.layers import nn_ops
label_desc: Use image node_id 2 for calling a global function with example usage: extract_tubelets_from_video(x, tubelet_size, select_central_frame) and returns: x