code: def pooling(
      x, pooling_mode="CLS", key_padding_mask=None, batch_first=True
  ):
      if pooling_mode not in ["MEAN", "SUM", "CLS", "LAST", "FLATTEN"]:
          raise NotImplementedError(
              f"pooling_mode must be MEAN, SUM, CLS, LAST, FLATTEN"
          )
      if pooling_mode in ["MEAN", "SUM"]:
          if key_padding_mask is not None:
              mask = rearrange(
                  ~key_padding_mask.bool_matrix,
                  "b s -> b s 1" if batch_first else "b s -> s b 1",
              )
              x = x.masked_fill(mask, 0)
          s = reduce(
              x,
              "b s ... -> b ..." if batch_first else "s b ... -> b ...",
              "sum",
          )
          if pooling_mode == "SUM":
              return s
          else:
              if key_padding_mask is None:
                  return s / x.shape[1 if batch_first else 0]
              else:
                  lengths = rearrange(
                      key_padding_mask._lengths, "b -> b 1"
                  )
                  return s / lengths
      elif pooling_mode == "CLS":
          return x[:, 0] if batch_first else x[0]
      elif pooling_mode == "LAST":
          if key_padding_mask is None:
              return x[:, -1] if batch_first else x[-1]
          else:
              lengths = key_padding_mask._lengths
              if batch_first:
                  batch_size = x.shape[0]
                  return x[
                      torch.arange(batch_size, device=x.device),
                      lengths - 1,
                  ]
              else:
                  batch_size = x.shape[1]
                  return x[
                      lengths - 1,
                      torch.arange(batch_size, device=x.device),
                  ]
      elif pooling_mode == "FLATTEN":
          return rearrange(
              x,
              "b ... -> b (...)"
              if batch_first
              else "s b ... -> b (s ...)",
          )
  
imports: import math
  from functools import partial
  from collections import namedtuple
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.modules.utils import _pair
  import hydra
  from einops import reduce, rearrange
label_desc: Use image node_id 1 for calling a global function with example usage: pooling(x, pooling_mode, key_padding_mask, batch_first) and returns: s, unknown, unknown, unknown, unknown, rearrange, x, x