Project import generated by Copybara.

PiperOrigin-RevId: 200715809

Author: tf-transform-team

RELEASE.md

@@ -28,6 +28,9 @@
 * tft.mean now supports SparseTensor when reduce_instance_dimensions=True.
   In this case it returns a scalar mean computed over the non-missing values of
   the SparseTensor.
+* tft.mean now supports SparseTensor when reduce_instance_dimensions=False.
+  In this case it returns a vector mean computed over the non-missing values of
+  the SparseTensor.
 * Update examples to use "core" TensorFlow estimator API (`tf.estimator`).
 
 ## Breaking changes

tensorflow_transform/analyzers.py

@@ -470,16 +470,23 @@ def mean(x, reduce_instance_dims=True, name=None, output_dtype=None):
     if output_dtype is None:
       raise TypeError('Tensor type %r is not supported' % x.dtype)
   sum_dtype, sum_fn = _sum_combine_fn_and_dtype(x.dtype)
-  if isinstance(x, tf.SparseTensor):
-    if not reduce_instance_dims:
-      raise TypeError(
-          'SparseTensor is only supported when reduce_instance_dims=True')
-    x = x.values
   with tf.name_scope(name, 'mean'):
-    # For now _numeric_combine will return a tuple with as many elements as the
-    # input tuple.
+    if isinstance(x, tf.SparseTensor):
+      if reduce_instance_dims:
+        ones_values, x_values = tf.ones_like(x.values), x.values
+      else:
+        sparse_ones = tf.SparseTensor(
+            indices=x.indices,
+            values=tf.ones_like(x.values),
+            dense_shape=x.dense_shape)
+        ones_values = tf.sparse_reduce_sum(sparse_ones, axis=0, keep_dims=True)
+        x = tf.cast(x, output_dtype)
+        ones_values = tf.cast(ones_values, output_dtype)
+        x_values = tf.sparse_reduce_sum(x, axis=0, keep_dims=True)
+    else:
+      ones_values, x_values = tf.ones_like(x), x
     x_count, x_sum = _numeric_combine(  # pylint: disable=unbalanced-tuple-unpacking
-        [tf.ones_like(x), x],
+        [ones_values, x_values],
         sum_fn,
         reduce_instance_dims,
         output_dtypes=[sum_dtype, sum_dtype])

tensorflow_transform/api.py

@@ -48,8 +48,6 @@ def preprocessing_fn(inputs):
 Beam implementation.
 """
 
-import collections
-
 import tensorflow as tf
 from tensorflow_transform import analyzers
 
@@ -168,37 +166,32 @@ def _convert_label(x):
   return FunctionApplication(fn, args).user_output
 
 
-# min_value and max_value are tensor names.
-_SchemaOverride = collections.namedtuple(
-    'SchemaOverride', ['min_value', 'max_value'])
-
-
-_TF_METADATA_TENSORS_COLLECTION = 'tft_metadata_tensors'
-_TF_METADATA_SCHEMA_OVERRIDES_COLLECTION = 'tft_metadata_schema_overrides'
+# Names of collections, which should all be the same length and contain tensors.
+# Each tensor in the first collection should have its min/max described by the
+# tensors in the other two collections.
+_TF_METADATA_TENSOR_COLLECTION = 'tft_schema_override_tensor'
+_TF_METADATA_TENSOR_MIN_COLLECTION = 'tft_schema_override_min'
+_TF_METADATA_TENSOR_MAX_COLLECTION = 'tft_schema_override_max'
 
 
 def set_tensor_schema_overrides(tensor, min_value, max_value):
-  """Override parts of the schema of a `Tensor` or `SparseTensor`."""
-  if not (isinstance(tensor, tf.Tensor) or isinstance(tensor, tf.SparseTensor)):
-    raise ValueError(
-        'tensor {} was not a Tensor or SparseTensor'.format(tensor))
+  """Override parts of the schema of a `Tensor`."""
+  if not isinstance(tensor, tf.Tensor):
+    raise ValueError('tensor {} was not a Tensor'.format(tensor))
   if not isinstance(min_value, tf.Tensor):
     raise ValueError('min_vaue {} was not a Tensor'.format(min_value))
   if not isinstance(max_value, tf.Tensor):
     raise ValueError('max_vaue {} was not a Tensor'.format(min_value))
-
-  tf.add_to_collection(_TF_METADATA_TENSORS_COLLECTION, tensor)
-
-  # Construct a _SchemaOverride using the tensor names of min_value and
-  # max_value.
-  tf.add_to_collection(_TF_METADATA_SCHEMA_OVERRIDES_COLLECTION,
-                       _SchemaOverride(min_value.name, max_value.name))
+  tf.add_to_collection(_TF_METADATA_TENSOR_COLLECTION, tensor)
+  tf.add_to_collection(_TF_METADATA_TENSOR_MIN_COLLECTION, min_value)
+  tf.add_to_collection(_TF_METADATA_TENSOR_MAX_COLLECTION, max_value)
 
 
 def get_tensor_schema_overrides():
-  """Gets a dict from `Tensor` or `SparseTensor`s to `_SchemaOverride`s."""
-  tensors = tf.get_collection(_TF_METADATA_TENSORS_COLLECTION)
-  schema_overrides = tf.get_collection(_TF_METADATA_SCHEMA_OVERRIDES_COLLECTION)
-  assert len(tensors) == len(schema_overrides), '{} != {}'.format(
-      tensors, schema_overrides)
-  return dict(zip(tensors, schema_overrides))
+  """Gets a dict from `Tensor`s to pairs of `Tensor`s containing min/max."""
+  tensors = tf.get_collection(_TF_METADATA_TENSOR_COLLECTION)
+  min_values = tf.get_collection(_TF_METADATA_TENSOR_MIN_COLLECTION)
+  max_values = tf.get_collection(_TF_METADATA_TENSOR_MAX_COLLECTION)
+  assert len(tensors) == len(min_values), '{} != {}'.format(tensors, min_values)
+  assert len(tensors) == len(max_values), '{} != {}'.format(tensors, max_values)
+  return dict(zip(tensors, zip(min_values, max_values)))

tensorflow_transform/beam/impl.py

@@ -408,8 +408,6 @@ def _write_saved_transform(graph, inputs, outputs, saved_model_dir):
     # warnings.
     # pylint: disable=protected-access
     collections_blacklist = [
-        tft_api._TF_METADATA_TENSORS_COLLECTION,
-        tft_api._TF_METADATA_SCHEMA_OVERRIDES_COLLECTION,
         tft_api.FUNCTION_APPLICATION_COLLECTION,
         tft_analyzers.ANALYZER_COLLECTION
     ]
@@ -629,112 +627,60 @@ def expand(self, inputs):
     return result
 
 
-class _ComputeDeferredMetadata(beam.PTransform):
-  """Extracts values of tensors from a transform function.
+def _augment_metadata(saved_model_dir, metadata):
+  """Augments the metadata with min/max values stored in the SavedModel.
 
-  This transform takes the path to a SavedModel in its constructor, and in its
-  expand() method accepts a mapping from tensors to PCollections.  When run, it
-  replaces the tensors corresponding to the keys of this mapping, with the
-  values wrapped in the PCollections.  It then extracts the values of some
-  tensors in the new graph.  This allows us to compute values that depend on
-  values in the tensor-PCollection mapping in arbitrary ways, where the values
-  are represented by tensors in the graph that depend on the tensor-PCollection
-  mapping (but not on the inputs to the graph).
+  Takes the min/max values of tensors stored in the SavedModel, and uses these
+  to augment the metadata.  For each feature in the metadata, the min/max of
+  the corresponding `Tensor` are used to augment the schema.  For a feature
+  represented by a `SparseTensor` we use the min/max for the `values` field of
+  the `SparseTensor`.
 
   Args:
+    saved_model_dir: Location of a SavedModel
     metadata: A `DatasetMetadata`
-    column_schema_overrides: A dict from column names to `api._SchemaOverride`s
-    saved_model_dir: The model to extract the constants from.
-    pipeline: The beam Pipeline.
-  """
-
-  def __init__(self, metadata, column_schema_overrides, saved_model_dir,
-               pipeline):
-    self._metadata = metadata
-    self._column_schema_overrides = column_schema_overrides
-    self._saved_model_dir = saved_model_dir
-    # Generally the pipeline is inferred from its inputs, however we need
-    # to know the pipeline for beam.Create.
-    self.pipeline = pipeline
 
-  def expand(self, tensor_pcoll_mapping):
-    """Converts a dict of statistics to a transform function.
-
-    Args:
-      tensor_pcoll_mapping: A dictionary mapping `Tensor`s to a singleton
-          PCollection containing a _TensorValue.
-
-    Returns:
-      A dict from tensor names to singleton `PCollection`s.
-    """
-    # Convert tensor_value_mapping into a DictPCollectionView so it can be
-    # passed as a side input to the beam Map below.
-    tensor_value_pairs = []
-    for name, pcoll in six.iteritems(tensor_pcoll_mapping):
-      tensor_value_pairs.append(
-          pcoll
-          | 'AddName[%s]' % name >> beam.Map(lambda x, name=name: (name, x)))
-    tensor_value_mapping = beam.pvalue.AsDict(
-        tensor_value_pairs
-        | 'MergeTensorValuePairs' >> beam.Flatten(pipeline=self.pipeline))
-
-    def compute_deferred_metadata(metadata, column_schema_overrides,
-                                  saved_model_dir, tensor_value_mapping):
-      """Extracts constant values from graph."""
-      tensor_names = {
-          tensor_name
-          for override in six.itervalues(column_schema_overrides)
-          for tensor_name in [override.min_value, override.max_value]}
-
-      graph = tf.Graph()
-      with graph.as_default():
-        tensor_replacement_map = {}
-        for orig_tensor_name, (value,
-                               is_asset) in six.iteritems(tensor_value_mapping):
-          new_tensor = tf.constant(value)
-          if is_asset:
-            # Any newly frozen constant tensors containing filenames must be
-            # added to the ASSET_FILENAMES collection.
-            graph.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, new_tensor)
-          tensor_replacement_map[orig_tensor_name] = new_tensor
-
-        with tf.Session(graph=graph) as session:
-          tensors_by_name = (
-              saved_transform_io.fetch_tensor_values(
-                  saved_model_dir, tensor_replacement_map, tensor_names))
-          session.run(tf.global_variables_initializer())
-          session.run(tf.tables_initializer())
-          tensor_values_by_name = session.run(tensors_by_name)
-
-      new_column_schemas = {}
-      for key, column_schema in six.iteritems(metadata.schema.column_schemas):
-        if key in column_schema_overrides:
-          override = column_schema_overrides[key]
-          min_value = tensor_values_by_name[override.min_value]
-          max_value = tensor_values_by_name[override.max_value]
-          assert column_schema.domain.dtype == tf.int64
-          assert isinstance(column_schema.domain, dataset_schema.IntDomain)
-          # Create a new column schema.  An override always results in a
-          # categorical column.
-          new_column_schemas[key] = dataset_schema.ColumnSchema(
-              dataset_schema.IntDomain(tf.int64, min_value, max_value,
-                                       is_categorical=True),
-              column_schema.axes,
-              column_schema.representation)
-        else:
-          new_column_schemas[key] = column_schema
-
-      return dataset_metadata.DatasetMetadata(dataset_schema.Schema(
-          new_column_schemas))
-
-    return (
-        self.pipeline
-        | 'CreateMetadata' >> beam.Create([self._metadata])
-        | 'ExtractScalarConstants' >> beam.Map(
-            compute_deferred_metadata,
-            column_schema_overrides=self._column_schema_overrides,
-            saved_model_dir=self._saved_model_dir,
-            tensor_value_mapping=tensor_value_mapping))
+  Returns:
+    An augmented DatasetMetadata.  The original DatasetMetadata is unchanged.
+  """
+  with tf.Graph().as_default() as graph:
+    with tf.Session(graph=graph) as session:
+      _, output_tensor_by_name = (
+          saved_transform_io.partially_apply_saved_transform_internal(
+              saved_model_dir, {}))
+
+      # Get overrides for the min/max of tensors from the graph, and use these
+      # determine overrides for the min/max of the outputs of the graph.
+      tensor_schema_overrides = tft_api.get_tensor_schema_overrides()
+      column_schema_overrides = {}
+      for name, tensor in six.iteritems(output_tensor_by_name):
+        if isinstance(tensor, tf.SparseTensor):
+          tensor = tensor.values
+        if tensor in tensor_schema_overrides:
+          column_schema_overrides[name] = tensor_schema_overrides[tensor]
+
+      session.run(tf.global_variables_initializer())
+      session.run(tf.tables_initializer())
+      column_schema_override_values = session.run(column_schema_overrides)
+
+  new_column_schemas = {}
+  for key, column_schema in six.iteritems(metadata.schema.column_schemas):
+    if key in column_schema_override_values:
+      min_value, max_value = column_schema_override_values[key]
+      assert column_schema.domain.dtype == tf.int64
+      assert isinstance(column_schema.domain, dataset_schema.IntDomain)
+      # Create a new column schema.  An override always results in a
+      # categorical column.
+      new_column_schemas[key] = dataset_schema.ColumnSchema(
+          dataset_schema.IntDomain(tf.int64, min_value, max_value,
+                                   is_categorical=True),
+          column_schema.axes,
+          column_schema.representation)
+    else:
+      new_column_schemas[key] = column_schema
+
+  return dataset_metadata.DatasetMetadata(dataset_schema.Schema(
+      new_column_schemas))
 
 
 class AnalyzeDataset(beam.PTransform):
@@ -860,23 +806,15 @@ def expand(self, dataset):
       # refer to values of tensors in the graph.  The override tensors must
       # be "constant" in that they don't depend on input data.  The tensors can
       # depend on analyzer outputs though.  This allows us to set metadata that
-      # depends on analyzer outputs. _ComputeDeferredMetadata will use
-      # tensor_pcoll_mapping to compute the metadata in a deferred manner, once
-      # the analyzer outputs are known.
+      # depends on analyzer outputs. _augment_metadata will use the analyzer
+      # outputs stored in `transform_fn` to compute the metadata in a
+      # deferred manner, once the analyzer outputs are known.
       metadata = dataset_metadata.DatasetMetadata(
           schema=impl_helper.infer_feature_schema(outputs))
 
-      tensor_schema_overrides = tft_api.get_tensor_schema_overrides()
-      column_schema_overrides = {
-          key: tensor_schema_overrides[tensor]
-          for key, tensor in six.iteritems(outputs)
-          if tensor in tensor_schema_overrides}
-
       deferred_metadata = (
-          tensor_pcoll_mapping
-          | 'ComputeDeferredMetadata' >>
-          _ComputeDeferredMetadata(metadata, column_schema_overrides,
-                                   saved_model_dir, input_values.pipeline))
+          transform_fn
+          | 'ComputeDeferredMetadata' >> beam.Map(_augment_metadata, metadata))
 
       full_metadata = beam_metadata_io.BeamDatasetMetadata(
           metadata, deferred_metadata)

tensorflow_transform/beam/impl_test.py

@@ -1282,7 +1282,7 @@ def analyzer_fn(inputs):
     self.assertAnalyzerOutputs(
         input_data, input_metadata, analyzer_fn, expected_outputs)
 
-  def testNumericMeanWithSparseTensor(self):
+  def testNumericMeanWithSparseTensorReduceTrue(self):
 
     def analyzer_fn(inputs):
       return {'mean': tft.mean(inputs['a'])}
@@ -1295,6 +1295,52 @@ def analyzer_fn(inputs):
     self.assertAnalyzerOutputs(input_data, input_metadata, analyzer_fn,
                                expected_outputs)
 
+  def testNumericMeanWithSparseTensorReduceFalse(self):
+
+    def analyzer_fn(inputs):
+      return {'mean': tft.mean(inputs['sparse'], False)}
+
+    input_data = [{
+        'sparse': ([0, 1], [0., 1.])
+    }, {
+        'sparse': ([1, 3], [2., 3.])
+    }]
+    input_metadata = dataset_metadata.DatasetMetadata({
+        'sparse':
+            sch.ColumnSchema(
+                tf.float32, [4],
+                sch.SparseColumnRepresentation(
+                    'val', [sch.SparseIndexField('idx', False)]))
+    })
+    expected_outputs = {
+        'mean': np.array([0., 1.5, float('nan'), 3.], np.float32)
+    }
+    self.assertAnalyzerOutputs(input_data, input_metadata, analyzer_fn,
+                               expected_outputs)
+
+  def testNumericMeanWithSparseTensorReduceFalseOverflow(self):
+
+    def analyzer_fn(inputs):
+      return {'mean': tft.mean(inputs['sparse'], False)}
+
+    input_data = [{
+        'sparse': ([0, 1], [1, 1])
+    }, {
+        'sparse': ([1, 3], [2147483647, 3])
+    }]
+    input_metadata = dataset_metadata.DatasetMetadata({
+        'sparse':
+            sch.ColumnSchema(
+                tf.int32, [4],
+                sch.SparseColumnRepresentation(
+                    'val', [sch.SparseIndexField('idx', False)]))
+    })
+    expected_outputs = {
+        'mean': np.array([1., 1073741824., float('nan'), 3.], np.float32)
+    }
+    self.assertAnalyzerOutputs(input_data, input_metadata, analyzer_fn,
+                               expected_outputs)
+
   def testNumericAnalyzersWithSparseInputs(self):
     def repeat(in_tensor, value):
       batch_size = tf.shape(in_tensor)[0]
@@ -1327,15 +1373,6 @@ def size_fn(inputs):
           return {'size': repeat(inputs['a'], tft.size(inputs['a']))}
         _ = input_dataset | beam_impl.AnalyzeDataset(size_fn)
 
-      with self.assertRaises(TypeError):
-        def mean_fn(inputs):
-          return {
-              'mean':
-                  repeat(inputs['a'],
-                         tft.mean(inputs['a'], reduce_instance_dims=False))
-          }
-        _ = input_dataset | beam_impl.AnalyzeDataset(mean_fn)
-
       with self.assertRaises(TypeError):
         def var_fn(inputs):
           return {'var': repeat(inputs['a'], tft.var(inputs['a']))}