def dense_layer_2d(input_tensor,
                   output_size,
                   initializer,
                   activation,
                   name=None):
  """A dense layer with 2D kernel.

  Args:
    input_tensor: Float tensor with rank 3.
    output_size: The size of output dimension.
    initializer: Kernel initializer.
    activation: Actication function.
    name: The name scope of this layer.

  Returns:
    float logits Tensor.
  """
  last_dim = get_shape_list(input_tensor)[-1]
  with tf.variable_scope(name):
    w = tf.get_variable(
        name="kernel", shape=[last_dim, output_size], initializer=initializer)
    b = tf.get_variable(
        name="bias", shape=[output_size], initializer=tf.zeros_initializer)

  ret = tf.einsum("abc,cd->abd", input_tensor, w)
  ret += b
  if activation is not None:
    return activation(ret)
  else:
    return ret
A dense layer with 2D kernel.

  Args:
    input_tensor: Float tensor with rank 3.
    output_size: The size of output dimension.
    initializer: Kernel initializer.
    activation: Actication function.
    name: The name scope of this layer.

  Returns:
    float logits Tensor.
  
language/canine/bert_modeling.py