def _post_process_regress()

in tfx_bsl/beam/run_inference.py [0:0]

• 30 lines of code
• 10 McCabe index (conditional complexity)

def _post_process_regress(
    examples: List[tf.train.Example],
    outputs: Mapping[Text, np.ndarray]) -> List[regression_pb2.Regression]:
  """Returns regressions from inference output."""

  if tf.saved_model.REGRESS_OUTPUTS not in outputs:
    raise ValueError('No regression outputs found in outputs: %s' %
                     outputs.keys())
  output = outputs[tf.saved_model.REGRESS_OUTPUTS]
  batch_size = len(examples)
  if not (output.ndim == 1 or (output.ndim == 2 and output.shape[1] == 1)):
    raise ValueError("""Expected output Tensor shape to be either [batch_size]
                     or [batch_size, 1] but got %s""" % output.shape)
  if batch_size != output.shape[0]:
    raise ValueError(
        'Input batch size did not match output batch size: %s vs %s' %
        (batch_size, output.shape[0]))
  if output.dtype != tf.float32.as_numpy_dtype:
    raise ValueError('Expected output Tensor of %s. Got: %s' %
                     (tf.float32.as_numpy_dtype, output.dtype))
  if output.size != batch_size:
    raise ValueError('Expected output batch size to be %s. Got: %s' %
                     (batch_size, output.size))
  flatten_output = output.flatten()
  result = []
  for value in flatten_output:
    regression = regression_pb2.Regression()
    regression.value = value
    result.append(regression)
  # Add additional check to save downstream consumer checks.
  if len(result) != len(examples):
    raise RuntimeError('Regression length does not match examples')
  return result