def head_projection(h, d_model, n_head, d_head, kernel_initializer, name):
  """Project hidden states to a specific head with a 4D-shape."""
  proj_weight = tf.get_variable('{}/kernel'.format(name),
                                [d_model, n_head, d_head], dtype=h.dtype,
                                initializer=kernel_initializer)
  head = tf.einsum('ibh,hnd->ibnd', h, proj_weight)

  return head
Project hidden states to a specific head with a 4D-shape.
src/modeling.py