code: def pooling(mode: Pooling):
      def pool_cls(inp):
          return inp[:, 0, :]
  
      def pool_mean(inp):
          return inp.mean(dim=1)
  
      return {Pooling.MEAN: pool_mean, Pooling.CLS: pool_cls}[mode]
  
imports: from enum import Enum
  from typing import Dict, Union
  import pytorch_lightning
  import torch
  import torch.nn
  from xformers.components import build_attention
  from xformers.components.multi_head_dispatch import MultiHeadDispatchConfig
  from xformers.factory import xFormer, xFormerConfig, xFormerEncoderConfig
  from xformers.utils import generate_matching_config
label_desc: Use image node_id 1 for calling a global function with example usage: pooling(mode) and returns: dict, inp, inp