code: def update_input_mask(
      input_mask: Optional[jnp.ndarray],
      taped_x: jnp.ndarray,
      ada_tt_len: Optional[jnp.ndarray],
      ac_config: ml_collections.ConfigDict,
  ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]:
      """Updates input mask and taped_x if needed."""
      logging_input_mask = None
      bs, in_len, _ = taped_x.shape
      num_tape_tokens = ac_config.num_tape_tokens
      if ac_config.get("dynamic_tape_length"):
          if input_mask is None:
              # Set the mask to one for all the tokens from the original input.
              input_mask = jnp.ones((bs, (in_len - num_tape_tokens)))
          assert ada_tt_len is not None
          tape_mask = (
              jnp.tile(jnp.arange(num_tape_tokens), (bs, 1))
              < ada_tt_len[..., None]
          )
          input_mask = jnp.concatenate(
              [input_mask, tape_mask.astype(input_mask.dtype)], axis=1
          )
          logging_input_mask = input_mask
      elif (
          input_mask is not None
          # Only update the mask if its size doesn't match the current input,
          # which means we need mask to also cover added tape tokens.
          and in_len != input_mask.shape[1]
      ):
          new_len = taped_x.shape[1]
          # Update the input mask to include tape tokens.
          tape_mask = jnp.ones(
              (input_mask.shape[0], (new_len - input_mask.shape[1]))
          )
          input_mask = jnp.concatenate([input_mask, tape_mask], axis=1)
  
      return input_mask, logging_input_mask, taped_x
  
imports: from typing import Any, Callable, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 2 for calling a global function with example usage: update_input_mask(input_mask, taped_x, ada_tt_len, ac_config) and returns: input_mask,  logging_input_mask,  taped_x