code: def prepare_encoder_input(
      inputs: Sequence[jnp.ndarray],
      masks: Sequence[jnp.ndarray],
      pos_embeds: Sequence[jnp.ndarray],
      level_embeds: jnp.ndarray,
  ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
      """Flatten all image dimensions and add level embed to position embed."""
      input_flattened = []
      mask_flattened = []
      level_pos_embed_flattened = []
      for inp, mask, pos_embed, level_embed in zip(
          inputs, masks, pos_embeds, level_embeds
      ):
          bs, h, w, c = inp.shape
          inp = inp.reshape(bs, h * w, c)
          mask = mask.reshape(bs, h * w)
          pos_embed = pos_embed.reshape(bs, h * w, c)
          lvl_pos_embed = pos_embed + level_embed
  
          level_pos_embed_flattened.append(lvl_pos_embed)
          input_flattened.append(inp)
          mask_flattened.append(mask)
  
      input_flattened = jnp.concatenate(input_flattened, 1)
      mask_flattened = jnp.concatenate(mask_flattened, 1)
      level_pos_embed_flattened = jnp.concatenate(
          level_pos_embed_flattened, 1
      )
      valid_ratios = jnp.stack(
          [get_mask_valid_ratio(m) for m in masks], 1
      )
      return (
          input_flattened,
          valid_ratios,
          level_pos_embed_flattened,
          mask_flattened,
      )
  
imports: import functools
  from typing import Sequence, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.deformable_detr.attention import MultiScaleDeformableAttention
  from scenic.projects.baselines.detr.model import MultiHeadDotProductAttention
label_desc: Use image node_id 4 for calling a global function with example usage: prepare_encoder_input(inputs, masks, pos_embeds, level_embeds) and returns: input_flattened,  valid_ratios,  level_pos_embed_flattened,  mask_flattened