class_docstr: A stack of layers for FunnelTransformer.
class_docstr_tok: ["A","stack","of","layers","for","FunnelTransformer","."]
docstr: Update the length dimension, sequence_id and position information.
docstr_tok: ["Update","the","length","dimension",",","sequence_id","and","position","information","."]
code: def update_context(self, context, x, pool_dim_name):
      """Update the length dimension, sequence_id and position information."""
  
      pooled_seq_length = x.shape.get_dim_by_name(pool_dim_name).size
      # For position, we slice the first `pooled_seq_length` indices instead of
      # striding. This ensures that the 3rd position before the pooling becomes
      # 2nd position after pooling instead of remembering its position before
      # pooling.
      new_context_position = mtf.slice(
          context.position,
          begin=0,
          size=pooled_seq_length,
          slice_dim_name=pool_dim_name,
      )
      context.position = new_context_position
  
      pooled_seq_length = x.shape.get_dim_by_name(pool_dim_name).size
      new_length_dim = mtf.Dimension(
          name=pool_dim_name, size=pooled_seq_length
      )
  
      new_sequence_id = mtf.stride_tensor_1d(
          context.sequence_id,
          pool_dim=context.length_dim,
          pool_size=self.pooling_size,
      )
  
      context.length_dim = new_length_dim
      context.sequence_id = new_sequence_id
  
file_path: mesh/mesh_tensorflow/transformer/funnel_transformer.py
imports: import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 2 for calling the FunnelTransformerLayerStack obj's underlying member method code with example usage: obj.update_context(context, x, pool_dim_name) without return types