def get_feature_columns(tf_transform_output):
  """Returns the FeatureColumns for the model.

  Args:
    tf_transform_output: A `TFTransformOutput` object.

  Returns:
    A list of FeatureColumns.
  """
  feature_spec = tf_transform_output.transformed_feature_spec()
  # Wrap scalars as real valued columns.
  def get_shape(spec):
    if isinstance(spec, tf.io.SparseFeature):
      return spec.size
    return spec.shape

  return [
      tf.feature_column.numeric_column(key, shape=get_shape(feature_spec[key]))
      for key in (common.NUMERIC_FEATURE_KEYS + common.CATEGORICAL_FEATURE_KEYS)
  ]
Returns the FeatureColumns for the model.

  Args:
    tf_transform_output: A `TFTransformOutput` object.

  Returns:
    A list of FeatureColumns.
  
examples/census_example.py