code: def _column_stack(logits, loss):
      """Stacks logits and losses.
  
      In case that only one exists, returns that one.
      Args:
        logits: logits array
        loss: loss array
  
      Returns:
        stacked logits and losses (or only one if both do not exist).
      """
      if logits is None:
          return np.expand_dims(loss, axis=-1)
      if loss is None:
          return logits
      return np.column_stack((logits, loss))
  
imports: import contextlib
  import dataclasses
  import logging
  from typing import Optional
  import numpy
  from sklearn import ensemble
  from sklearn import linear_model
  from sklearn import model_selection
  from sklearn import neighbors
  from sklearn import neural_network
  from sklearn.utils import parallel_backend
  from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack import data_structures
label_desc: Use image node_id 3 for calling a global function with example usage: _column_stack(logits, loss) and returns: np, np, logits