code: def get_output_shapes(
      feature_target: str,
      patch_size: Patch,
      select_central_frame: Optional[bool] = None,
      channels: int = 3,
  ):
      """Returns the output shape, depending on the feature regression target."""
      if feature_target == FeatureTargets.RGB:
          if len(patch_size) == 3 and select_central_frame:
              output_elements = patch_size[:2] + (1, channels)
          else:
              output_elements = patch_size + (channels,)
          return np.prod(np.array(output_elements))
      elif feature_target == FeatureTargets.SPECTROGRAM:
          output_elements = patch_size + (channels,)
          return np.prod(np.array(output_elements))
      else:
          raise NotImplementedError(
              "Other feature targets not implemented yet."
          )
  
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 1 for calling a global function with example usage: get_output_shapes(feature_target, patch_size, select_central_frame, channels) and returns: np, np