from tensorflow.python.ops import array_ops

  """_TargetColumn for regression."""

  def __init__(self, loss_fn, label_name, weight_column_name, target_dimension):
    super(_RegressionTargetColumn, self).__init__(
        loss_fn=loss_fn,
        num_label_columns=target_dimension,
        label_name=label_name,
        weight_column_name=weight_column_name)

  def logits_to_predictions(self, logits, proba=False):
    if self.num_label_columns == 1:
      return array_ops.squeeze(logits, squeeze_dims=[1])
    return logits

  def get_eval_ops(self, features, logits, targets, metrics=None):
    loss = self.loss(logits, targets, features)
    result = {"loss": metrics_lib.streaming_mean(loss)}
    if metrics:
      predictions = self.logits_to_predictions(logits, proba=False)
      result.update(_run_metrics(predictions, targets, metrics,
                                 self.get_weight_tensor(features)))
    return result
