def __call__(self, x, mask=None):
    """Apply the Transformer pooler.

    Args:
      x: <float32>[batch_size, len, dim].
      mask: <int32>[batch_size, len].

    Returns:
      pooled: <float32>[batch_size, hidden_size].
    """
    if mask is None:
      mask = jnp.ones(x.shape[:2], dtype=jnp.int32)
    mask = mask[:, :, None]
    if self.pooling_operator == 'mean':
      masked_x = x * mask
      pooled = masked_x.sum(1) / jnp.maximum(mask.sum(1), 1e-9)
    elif self.pooling_operator == 'first':
      pooled = x[:, 0]
    elif self.pooling_operator == 'max':
      masked_x = x * mask + -1e10 * (1 - mask)
      pooled = jnp.max(x, axis=1)
    else:
      raise NotImplementedError
    pooled = nn.Dense(
        self.hidden_size,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        name='dense')(
            pooled)
    pooled = self.output_act(pooled)
    return pooled
Apply the Transformer pooler.

    Args:
      x: <float32>[batch_size, len, dim].
      mask: <int32>[batch_size, len].

    Returns:
      pooled: <float32>[batch_size, hidden_size].
    
language/gscan/xattn_model/model/layers.py