def last_dim_weighted_sum(x, weight_name, weight_init=None, keepdims=False):
  """Computes a weighted sum of the last dimension of `x`.

  Args:
    x: <float32>[..., x_dim]
    weight_name: Name of the weight variable to use
    weight_init: Initializer of the weight variable
    keepdims: Whether the output should hav an ending size one dim

  Returns:
    summed: <float32>[...] or <float32>[..., 1] iff `keepdims`
  """

  dim = x.shape.as_list()[-1]
  w = tf.get_variable(weight_name, dim, initializer=weight_init)
  out = tf.tensordot(x, w, [[len(x.shape) - 1], [0]])
  if keepdims:
    return tf.expand_dims(out, len(out.shape))
  else:
    return out
Computes a weighted sum of the last dimension of `x`.

  Args:
    x: <float32>[..., x_dim]
    weight_name: Name of the weight variable to use
    weight_init: Initializer of the weight variable
    keepdims: Whether the output should hav an ending size one dim

  Returns:
    summed: <float32>[...] or <float32>[..., 1] iff `keepdims`
  
language/boolq/utils/ops.py