Create the initial cache for Evolved Transformer fast decoding.

def _init_evolved_transformer_cache(cache, hparams, batch_size,
                                    attention_init_length, encoder_output,
                                    encoder_decoder_attention_bias,
                                    scope_prefix):
  """Create the initial cache for Evolved Transformer fast decoding."""
  key_channels = hparams.attention_key_channels or hparams.hidden_size
  value_channels = hparams.attention_value_channels or hparams.hidden_size
  num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
  vars_3d_num_heads = (
      hparams.num_heads if hparams.get("attention_variables_3d") else 0)

  # Add self-attentions.
  if cache is None:
    cache = {}
  cache.update({
      "layer_%d" % layer: {  # pylint: disable=g-complex-comprehension
          _SIXTEEN_HEAD_ATTENTION_NAME: {
              "k":
                  common_attention.split_heads(
                      tf.zeros(
                          [batch_size, attention_init_length, key_channels]),
                      _capped_double_heads(hparams.num_heads)),
              "v":
                  common_attention.split_heads(
                      tf.zeros(
                          [batch_size, attention_init_length, value_channels]),
                      _capped_double_heads(hparams.num_heads)),
          },
          _VANILLA_ATTENTION_NAME: {
              "k":
                  common_attention.split_heads(
                      tf.zeros(
                          [batch_size, attention_init_length, key_channels]),
                      hparams.num_heads),
              "v":
                  common_attention.split_heads(
                      tf.zeros(
                          [batch_size, attention_init_length, value_channels]),
                      hparams.num_heads),
          }
      } for layer in range(num_layers)
  })

  # Add branched layers. Pad with additional zeros for causal convolution.
  for layer in range(num_layers):
    cache["layer_%d" % layer][_CONV_BRANCHES_FIRST_LAYER_NAME] = tf.zeros([
        batch_size, attention_init_length + _DECODER_LEFT_CONV_PADDING,
        hparams.hidden_size
    ])
    cache["layer_%d" % layer][_CONV_BRANCHES_SECOND_LAYER_NAME] = tf.zeros([
        batch_size, attention_init_length + _DECODER_FINAL_CONV_PADDING,
        hparams.hidden_size * 2
    ])

  # Add encoder embedding attentions.
  if encoder_output is not None:
    cache = _add_attend_to_encoder_cache(
        cache=cache,
        attention_name=_FIRST_ATTEND_TO_ENCODER_NAME,
        hparams=hparams,
        num_layers=num_layers,
        key_channels=key_channels,
        value_channels=value_channels,
        vars_3d_num_heads=vars_3d_num_heads,
        scope_prefix=scope_prefix,
        encoder_output=encoder_output)
    cache = _add_attend_to_encoder_cache(
        cache=cache,
        attention_name=_SECOND_ATTEND_TO_ENCODER_NAME,
        hparams=hparams,
        num_layers=num_layers,
        key_channels=key_channels,
        value_channels=value_channels,
        vars_3d_num_heads=vars_3d_num_heads,
        scope_prefix=scope_prefix,
        encoder_output=encoder_output)

    cache["encoder_output"] = encoder_output
    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias

  return cache