Helper: create positional encoding parameters.

def _positional_encoding_new_params(input_shape, rng, max_len=2048):  # pylint: disable=invalid-name
  """Helper: create positional encoding parameters."""
  del rng
  # Check if we are operating on chunked inputs by checking if the first
  # shape is a list/tuple of shapes (otherwise it's an int or numpy array).
  is_chunked = isinstance(input_shape[0], (list, tuple))
  feature_depth = input_shape[0][-1] if is_chunked else input_shape[-1]
  pe = onp.zeros((max_len, feature_depth), dtype=onp.float32)
  position = onp.arange(0, max_len)[:, onp.newaxis]
  div_term = onp.exp(
      onp.arange(0, feature_depth, 2) * -(onp.log(10000.0) / feature_depth))
  pe[:, 0::2] = onp.sin(position * div_term)
  pe[:, 1::2] = onp.cos(position * div_term)
  pe = pe[onp.newaxis, :, :]  # [1, max_len, feature_depth]
  return np.array(pe)