code: def LinearPooling(
      shorten_factor,
      d_model,
      *args,
      dropout=0.0,
      mode="train",
      **kwargs,
  ):
      del args, kwargs
  
      return cb.Serial(
          core.Fn(
              "Shorten",
              lambda x: jnp.reshape(  # pylint: disable=g-long-lambda
                  # Shorten -- move to depth.  # pylint: disable=g-long-lambda
                  x,
                  (x.shape[0], x.shape[1] // shorten_factor, -1),
              ),
              n_out=1,
          ),
          core.Dense(d_model),
          core.Dropout(rate=dropout, mode=mode),
      )
  
file_path: trax/trax/layers/research/resampling.py
imports: from trax.fastmath import numpy
  from trax.layers import combinators
  from trax.layers import core
  from trax.layers.normalization import LayerNorm
  from trax.layers.pooling import AvgPool
  from trax.layers.research.rel_attention import RelativeAttentionLMLayer
label_desc: Use image node_id 2 for calling a global function with example usage: LinearPooling(shorten_factor, d_model) and returns: cb