from tensorflow.python.ops import array_ops

      return array_ops.reshape(
          math_ops.to_float(features[self._weight_column_name]), shape=(-1,))

  @property
  def problem_type(self):
    return self._problem_type

  def _weighted_loss(self, loss, weight_tensor):
    """Returns cumulative weighted loss."""
    unweighted_loss = array_ops.reshape(loss, shape=(-1,))
    weighted_loss = math_ops.multiply(unweighted_loss,
                                      array_ops.reshape(
                                          weight_tensor, shape=(-1,)))
    return weighted_loss

  def training_loss(self, logits, target, features, name="training_loss"):
    """Returns training loss tensor for this head.

    Training loss is different from the loss reported on the tensorboard as we
    should respect the example weights when computing the gradient.

      L = sum_{i} w_{i} * l_{i} / B
