(shorten_factor, d_model, is_upsampling, d_ff, n_heads,
                        dropout, dropout_shared_axes, mode, ff_activation,
                        context_bias_layer, location_bias_layer, total_pooling,
                        resampling_fn)
return [
      LayerNorm(),               # h
      cb.Branch(cb.Serial(
          resampling,
          LayerNorm(),
      ), None),                  # h', h
      cb.Serial(  # pylint: disable=g-long-ternary
          cb.Select([0, 2, 1, 2]),
          cb.Add(),
      ) if is_upsampling else [],
      cb.Residual(
          cb.Select([0, 1, 1]),  # h', h, h
          attention,
          _Dropout(),
      ),
      cb.Residual(
          LayerNorm(),
          feed_forward,
          _Dropout(),
      ),
  ]
Attention resampling.
Attention resampling.
Attention resampling .
def AttentionResampling(shorten_factor, d_model, is_upsampling, d_ff, n_heads,
                        dropout, dropout_shared_axes, mode, ff_activation,
                        context_bias_layer, location_bias_layer, total_pooling,
                        resampling_fn):
  """Attention resampling."""

  attention = RelativeAttentionLMLayer(
      d_model, context_bias_layer, location_bias_layer,
      total_pooling, n_heads=n_heads, dropout=dropout,
      mode=mode)

  feed_forward = FeedForwardBlock(
      d_model, d_ff, dropout, dropout_shared_axes, mode, ff_activation)

  resampling = resampling_fn(shorten_factor, d_model,
                             mode=mode)

  def _Dropout():
    return core.Dropout(rate=dropout, shared_axes=dropout_shared_axes,
                        mode=mode)

  return [
      LayerNorm(),               # h
      cb.Branch(cb.Serial(
          resampling,
          LayerNorm(),
      ), None),                  # h', h
      cb.Serial(  # pylint: disable=g-long-ternary
          cb.Select([0, 2, 1, 2]),
          cb.Add(),
      ) if is_upsampling else [],
      cb.Residual(
          cb.Select([0, 1, 1]),  # h', h, h
          attention,
          _Dropout(),
      ),
      cb.Residual(
          LayerNorm(),
          feed_forward,
          _Dropout(),
      ),
  ]
def AttentionResampling ( shorten_factor , d_model , is_upsampling , d_ff , n_heads , dropout , dropout_shared_axes , mode , ff_activation , context_bias_layer , location_bias_layer , total_pooling , resampling_fn ) : attention = RelativeAttentionLMLayer ( d_model , context_bias_layer , location_bias_layer , total_pooling , n_heads = n_heads , dropout = dropout , mode = mode ) feed_forward = FeedForwardBlock ( d_model , d_ff , dropout , dropout_shared_axes , mode , ff_activation ) resampling = resampling_fn ( shorten_factor , d_model , mode = mode ) def _Dropout ( ) : return core . Dropout ( rate = dropout , shared_axes = dropout_shared_axes , mode = mode ) return [ LayerNorm ( ) , # h cb . Branch ( cb . Serial ( resampling , LayerNorm ( ) , ) , None ) , # h', h cb . Serial ( # pylint: disable=g-long-ternary cb . Select ( [ 0 , 2 , 1 , 2 ] ) , cb . Add ( ) , ) if is_upsampling else [ ] , cb . Residual ( cb . Select ( [ 0 , 1 , 1 ] ) , # h', h, h attention , _Dropout ( ) , ) , cb . Residual ( LayerNorm ( ) , feed_forward , _Dropout ( ) , ) , ]