code: def apply_encoder(
      x_unmasked_dict: ArrayDict,
      encoder_strategy: str,
      mlp_dim: int,
      num_layers: int,
      num_heads: int,
      attention_config: ml_collections.ConfigDict,
      dropout_rate: float,
      attention_dropout_rate: float,
      stochastic_droplayer_rate: float,
      dtype: jnp.dtype,
      normalise_encoder_output: bool,
      train: bool,
      fusion_layers: Optional[int] = None,
  ) -> ArrayDict:
      """Apply the encoder for each modality."""
  
      def separate_encoders():
          x_unmasked_encoded_dict = {}
          for key, x_unmasked_ in x_unmasked_dict.items():
              x_unmasked_encoded = vivit_model.Encoder(
                  temporal_dims=None,
                  mlp_dim=mlp_dim,
                  num_layers=num_layers,
                  num_heads=num_heads,
                  attention_config=attention_config,
                  dropout_rate=dropout_rate,
                  attention_dropout_rate=attention_dropout_rate,
                  stochastic_droplayer_rate=stochastic_droplayer_rate,
                  dtype=dtype,
                  positional_embedding="none",  # Has already been added.
                  normalise_output=normalise_encoder_output,
                  name=f"Transformer_{key}",
              )(x_unmasked_, train=train)
  
              x_unmasked_encoded = nn_layers.IdentityLayer(
                  name=f"encoder_output_{key}"
              )(x_unmasked_encoded)
  
              x_unmasked_encoded_dict[key] = x_unmasked_encoded
          return x_unmasked_encoded_dict
  
      def concat_and_encode():
          x_unmasked_encoded_dict = {}
          encoder = vivit_model.Encoder(
              temporal_dims=None,
              mlp_dim=mlp_dim,
              num_layers=num_layers,
              num_heads=num_heads,
              attention_config=attention_config,
              dropout_rate=dropout_rate,
              attention_dropout_rate=attention_dropout_rate,
              stochastic_droplayer_rate=stochastic_droplayer_rate,
              dtype=dtype,
              positional_embedding="none",  # Has already been added.
              normalise_output=normalise_encoder_output,
              name="Transformer",
          )
  
          x_unmasked_input_list = []
          x_unmasked_key_list = []
          for key, x_unmasked_ in x_unmasked_dict.items():
              x_unmasked_input_list.append(x_unmasked_)
              x_unmasked_key_list.append(key)
  
          x_unmasked_concat_input = jnp.concatenate(
              x_unmasked_input_list, axis=1
          )
          x_unmasked_encoded_concat = encoder(
              x_unmasked_concat_input, train=train
          )
          start_index = 0
          for key, x_unmasked_ in zip(
              x_unmasked_key_list, x_unmasked_input_list
          ):
              end_index = start_index + x_unmasked_.shape[1]
              x_unmasked_encoded_dict[key] = nn_layers.IdentityLayer(
                  name=f"encoder_output_{key}"
              )(x_unmasked_encoded_concat[:, start_index:end_index])
              start_index = end_index
              assert (
                  x_unmasked_encoded_dict[key].shape[1]
                  == x_unmasked_.shape[1]
              )
  
          return x_unmasked_encoded_dict
  
      def same_encoder():
          x_unmasked_encoded_dict = {}
          encoder = vivit_model.Encoder(
              temporal_dims=None,
              mlp_dim=mlp_dim,
              num_layers=num_layers,
              num_heads=num_heads,
              attention_config=attention_config,
              dropout_rate=dropout_rate,
              attention_dropout_rate=attention_dropout_rate,
              stochastic_droplayer_rate=stochastic_droplayer_rate,
              dtype=dtype,
              positional_embedding="none",  # Has already been added.
              normalise_output=normalise_encoder_output,
              name="Transformer",
          )
  
          for key, x_unmasked_ in x_unmasked_dict.items():
              x_unmasked_encoded = encoder(x_unmasked_, train=train)
              x_unmasked_encoded = nn_layers.IdentityLayer(
                  name=f"encoder_output_{key}"
              )(x_unmasked_encoded)
              x_unmasked_encoded_dict[key] = x_unmasked_encoded
          return x_unmasked_encoded_dict
  
      def separate_encoders_and_concat():
          assert fusion_layers is not None
          num_layers_single = num_layers - fusion_layers
          num_layers_concat = fusion_layers
          assert num_layers_single + num_layers_concat == num_layers
  
          x_unmasked_encoded_dict = {}
  
          for key, x_unmasked_ in x_unmasked_dict.items():
              x_unmasked_encoded = vivit_model.Encoder(
                  temporal_dims=None,
                  mlp_dim=mlp_dim,
                  num_layers=num_layers_single,
                  num_heads=num_heads,
                  attention_config=attention_config,
                  dropout_rate=dropout_rate,
                  attention_dropout_rate=attention_dropout_rate,
                  stochastic_droplayer_rate=stochastic_droplayer_rate,
                  dtype=dtype,
                  positional_embedding="none",  # Has already been added.
                  normalise_output=normalise_encoder_output,
                  name=f"Transformer_{key}",
              )(x_unmasked_, train=train)
  
              x_unmasked_encoded_dict[key] = x_unmasked_encoded
  
          encoder_concat = vivit_model.Encoder(
              temporal_dims=None,
              mlp_dim=mlp_dim,
              num_layers=num_layers_concat,
              num_heads=num_heads,
              attention_config=attention_config,
              dropout_rate=dropout_rate,
              attention_dropout_rate=attention_dropout_rate,
              stochastic_droplayer_rate=stochastic_droplayer_rate,
              dtype=dtype,
              positional_embedding="none",
              normalise_output=normalise_encoder_output,
              name="Transformer_concat",
          )
  
          x_unmasked_input_list = []
          x_unmasked_key_list = []
          for (
              key,
              x_unmasked_encoded_,
          ) in x_unmasked_encoded_dict.items():
              x_unmasked_input_list.append(x_unmasked_encoded_)
              x_unmasked_key_list.append(key)
  
          x_unmasked_concat_input = jnp.concatenate(
              x_unmasked_input_list, axis=1
          )
          x_unmasked_encoded_concat = encoder_concat(
              x_unmasked_concat_input, train=train
          )
  
          x_unmasked_encoded_dict_out = {}
          start_index = 0
          for key, x_unmasked_ in zip(
              x_unmasked_key_list, x_unmasked_input_list
          ):
              end_index = start_index + x_unmasked_.shape[1]
              x_unmasked_encoded_dict_out[
                  key
              ] = nn_layers.IdentityLayer(name=f"encoder_output_{key}")(
                  x_unmasked_encoded_concat[:, start_index:end_index]
              )
              start_index = end_index
              assert (
                  x_unmasked_encoded_dict_out[key].shape[1]
                  == x_unmasked_.shape[1]
              )
  
          return x_unmasked_encoded_dict_out
  
      if encoder_strategy == "separate_encoders":
          return separate_encoders()
      elif encoder_strategy == "concat_and_encode":
          return concat_and_encode()
      elif encoder_strategy == "same_encoder":
          return same_encoder()
      elif encoder_strategy == "separate_encoders_and_concat":
          return separate_encoders_and_concat()
      else:
          raise ValueError(
              f"The encoder strategy {encoder_strategy} is not supported!"
          )
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.av_mae import base_model
  from scenic.projects.av_mae import model_utils
  from scenic.projects.mbt import model
  from scenic.projects.vivit import model
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 1 for calling a global function with example usage: apply_encoder(x_unmasked_dict, encoder_strategy, mlp_dim, num_layers, num_heads, attention_config, dropout_rate, attention_dropout_rate, stochastic_droplayer_rate, dtype, normalise_encoder_output, train, fusion_layers) and returns: x_unmasked_encoded_dict, x_unmasked_encoded_dict, x_unmasked_encoded_dict, x_unmasked_encoded_dict_out, separate_encoders, concat_and_encode, same_encoder, separate_encoders_and_concat