tft.quantiles(reduce_instance_dims=False) errors with columns that are all NaN-valued #256

import pprint
import tempfile

import tensorflow as tf
import tensorflow_transform as tft
import tensorflow_transform.beam as tft_beam
from tensorflow_transform.tf_metadata import dataset_metadata
from tensorflow_transform.tf_metadata import schema_utils


def main():
  def preprocessing_fn(inputs):
    """Preprocess input columns into transformed columns."""
    features = [
        inputs['no_nans'],
        inputs['some_nans'],
        inputs['all_nans']
    ]
    concatenated = tf.stack(features, axis=-1)
    medians = tft.quantiles(x=concatenated, num_buckets=2, epsilon=0.01, reduce_instance_dims=False)[:, 0]
    return {
        'normalized': concatenated / medians
    }

  nan = float('nan')
  data = [
      {'no_nans': 0.0, 'some_nans': 0.0, 'all_nans': nan},
      {'no_nans': 1.0, 'some_nans': nan, 'all_nans': nan},
      {'no_nans': 2.0, 'some_nans': 2.0, 'all_nans': nan},
      {'no_nans': 3.0, 'some_nans': nan, 'all_nans': nan},
      {'no_nans': 4.0, 'some_nans': 4.0, 'all_nans': nan},
      {'no_nans': 5.0, 'some_nans': nan, 'all_nans': nan},
      {'no_nans': 6.0, 'some_nans': 6.0, 'all_nans': nan},
      {'no_nans': 7.0, 'some_nans': nan, 'all_nans': nan},
      {'no_nans': 8.0, 'some_nans': 8.0, 'all_nans': nan},
      {'no_nans': 9.0, 'some_nans': nan, 'all_nans': nan},
  ]

  raw_data_metadata = dataset_metadata.DatasetMetadata(
      schema_utils.schema_from_feature_spec({
          'no_nans': tf.io.FixedLenFeature([], tf.float32),
          'some_nans': tf.io.FixedLenFeature([], tf.float32),
          'all_nans': tf.io.FixedLenFeature([], tf.float32),
      }))

  with tft_beam.Context(temp_dir=tempfile.mkdtemp()):
    transformed_dataset, transform_fn = (  # pylint: disable=unused-variable
        (data, raw_data_metadata) | tft_beam.AnalyzeAndTransformDataset(
            preprocessing_fn))

  transformed_data, transformed_metadata = transformed_dataset  # pylint: disable=unused-variable

  pprint.pprint(transformed_data)

if __name__ == '__main__':
  main()