diff --git a/README.md b/README.md
index 7e74d99fc5..ee880c0397 100644
--- a/README.md
+++ b/README.md
@@ -57,7 +57,8 @@ combinations may also work.
 
 |tensorflow-model-analysis  |tensorflow    |apache-beam[gcp]|
 |---------------------------|--------------|----------------|
-|GitHub master              |1.9           |2.5.0           |
+|GitHub master              |1.9           |2.6.0           |
+|0.9.1                      |1.9           |2.6.0           |
 |0.9.0                      |1.9           |2.5.0           |
 |0.6.0                      |1.6           |2.4.0           |
 
diff --git a/RELEASE.md b/RELEASE.md
index ee2b464e0a..76cd09fccf 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -1,3 +1,22 @@
+# Release 0.9.1
+
+## Major Features and Improvements
+
+## Bug fixes and other changes
+
+*   Requires pre-installed TensorFlow >=1.10,<2.
+*   Updated ExampleCount to use the batch dimension as the example count. It
+    also now tries a few fallbacks if none of the standard keys are found in the
+    predictions dictionary: the first key in sorted order in the predictions
+    dictionary, or failing that, the first key in sorted order in the labels
+    dictionary, or failing that, it defaults to zero.
+*   Fix bug where we were mutating an element in a DoFn - this is prohibited in
+    the Beam model and can cause subtle bugs.
+
+## Breaking changes
+
+## Deprecations
+
 # Release 0.9.0
 
 ## Major Features and Improvements
diff --git a/examples/chicago_taxi/README.md b/examples/chicago_taxi/README.md
index 4b99f4815c..746d644ec9 100644
--- a/examples/chicago_taxi/README.md
+++ b/examples/chicago_taxi/README.md
@@ -45,8 +45,8 @@ Development for this example will be isolated in a Python virtual environment.
 This allows us to experiment with different versions of dependencies.
 
 There are many ways to install `virtualenv`, see the
-[TensorFlow install guides](/install) for different platforms, but here are a
-couple:
+[TensorFlow install guides](https://www.tensorflow.org/install) for different
+platforms, but here are a couple:
 
 * For Linux:
 
diff --git a/setup.py b/setup.py
index 46c092b61d..c3708d4b2c 100644
--- a/setup.py
+++ b/setup.py
@@ -243,7 +243,7 @@ def run(self):
         'tensorflow_model_analysis/static/vulcanized_template.html',
     ]),],
     'install_requires': [
-        'apache-beam[gcp]>=2.5,<3',
+        'apache-beam[gcp]>=2.6,<3',
         'grpc-google-iam-v1==0.11.1',
         'numpy>=1.10,<2',
         'jupyter>=1.0,<2',
diff --git a/tensorflow_model_analysis/BUILD b/tensorflow_model_analysis/BUILD
index 9ad0b6cc7d..3753e68979 100644
--- a/tensorflow_model_analysis/BUILD
+++ b/tensorflow_model_analysis/BUILD
@@ -1,10 +1,30 @@
 # Description:
 #   Python SDK for the TensorFlow Model Analysis API.
 
+package(default_visibility = ["//tensorflow_model_analysis:__subpackages__"])
+
 licenses(["notice"])  # Apache 2.0
 
 exports_files(["LICENSE"])
 
+py_library(
+    name = "tensorflow_model_analysis",
+    srcs = [
+        "__init__.py",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":constants",
+        ":version",
+        "//tensorflow_model_analysis/api:model_eval_lib",
+        "//tensorflow_model_analysis/api:tfma_unit_non_testonly",
+        "//tensorflow_model_analysis/eval_saved_model:export",
+        "//tensorflow_model_analysis/eval_saved_model:exporter",
+        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics",
+        "//tensorflow_model_analysis/slicer",
+    ],
+)
+
 py_library(
     name = "constants",
     srcs = ["constants.py"],
diff --git a/tensorflow_model_analysis/api/BUILD b/tensorflow_model_analysis/api/BUILD
deleted file mode 100644
index 5d1672d6ac..0000000000
--- a/tensorflow_model_analysis/api/BUILD
+++ /dev/null
@@ -1,67 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-exports_files(["LICENSE"])
-
-py_library(
-    name = "model_eval_lib",
-    srcs = ["model_eval_lib.py"],
-    deps = [
-        "//tensorflow_model_analysis:constants",
-        "//tensorflow_model_analysis:expect_apache_beam_installed",  # b/73825929
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/api/impl:api_types",
-        "//tensorflow_model_analysis/api/impl:evaluate",
-        "//tensorflow_model_analysis/api/impl:serialization",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics",
-        "//tensorflow_model_analysis/slicer",
-    ],
-)
-
-py_test(
-    name = "model_eval_lib_test",
-    srcs = ["model_eval_lib_test.py"],
-    shard_count = 3,
-    deps = [
-        ":model_eval_lib",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:testutil",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:csv_linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:custom_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics",
-        "//tensorflow_model_analysis/proto:metrics_for_slice_py_pb2",
-        "//tensorflow_model_analysis/slicer",
-    ],
-)
-
-py_library(
-    name = "tfma_unit_non_testonly",
-    srcs = ["tfma_unit.py"],
-    deps = [
-        "//tensorflow_model_analysis:expect_apache_beam_installed",
-        "//tensorflow_model_analysis/api/impl:evaluate",
-        "//tensorflow_model_analysis/eval_saved_model:load",
-        "//tensorflow_model_analysis/eval_saved_model:testutil_non_testonly",
-        "//tensorflow_model_analysis/slicer",
-    ],
-)
-
-py_library(
-    name = "tfma_unit",
-    testonly = 1,
-    deps = ["tfma_unit_non_testonly"],
-)
-
-py_test(
-    name = "tfma_unit_test",
-    srcs = ["tfma_unit_test.py"],
-    deps = [
-        ":tfma_unit",
-        "//tensorflow_model_analysis:expect_apache_beam_installed",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_estimator_extra_fields",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics",
-    ],
-)
diff --git a/tensorflow_model_analysis/api/impl/BUILD b/tensorflow_model_analysis/api/impl/BUILD
deleted file mode 100644
index a75e1bbb9a..0000000000
--- a/tensorflow_model_analysis/api/impl/BUILD
+++ /dev/null
@@ -1,74 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-exports_files(["LICENSE"])
-
-py_library(
-    name = "evaluate",
-    srcs = ["evaluate.py"],
-    deps = [
-        "//tensorflow_model_analysis:expect_apache_beam_installed",  # b/73825929
-        "//tensorflow_model_analysis:expect_six_installed",
-        "//tensorflow_model_analysis:expect_tensorflow_transform_installed",
-        "//tensorflow_model_analysis:types",
-        "//tensorflow_model_analysis/eval_saved_model:load",
-        "//tensorflow_model_analysis/eval_saved_model:util",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics:metric_keys",
-        "//tensorflow_model_analysis/extractors:feature_extractor",
-        "//tensorflow_model_analysis/proto:metrics_for_slice_py_pb2",
-        "//tensorflow_model_analysis/slicer",
-    ],
-)
-
-py_test(
-    name = "evaluate_test",
-    srcs = ["evaluate_test.py"],
-    shard_count = 4,
-    deps = [
-        ":evaluate",
-        "//tensorflow_model_analysis:expect_apache_beam_installed",  # b/73825929
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:testutil",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics:metric_keys",
-        "//tensorflow_model_analysis/proto:metrics_for_slice_py_pb2",
-        "//third_party/tensorflow/core:protos_all_py_pb2",
-    ],
-)
-
-py_library(
-    name = "serialization",
-    srcs = ["serialization.py"],
-    deps = [
-        ":api_types",
-        "//tensorflow_model_analysis:expect_apache_beam_installed",  # b/73825929
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:version",
-        "//tensorflow_model_analysis/api/impl:evaluate",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics",
-        "//tensorflow_model_analysis/proto:metrics_for_slice_py_pb2",
-        "//tensorflow_model_analysis/slicer",
-        "//third_party/py/numpy",
-    ],
-)
-
-py_test(
-    name = "serialization_test",
-    srcs = ["serialization_test.py"],
-    deps = [
-        ":serialization",
-        "//tensorflow_model_analysis:types",
-        "//tensorflow_model_analysis/eval_saved_model:testutil",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics:metric_keys",
-    ],
-)
-
-py_library(
-    name = "api_types",
-    srcs = ["api_types.py"],
-    deps = [
-        "//tensorflow_model_analysis:constants",
-        "//tensorflow_model_analysis/slicer",
-    ],
-)
diff --git a/tensorflow_model_analysis/api/impl/evaluate.py b/tensorflow_model_analysis/api/impl/evaluate.py
index 6fd69c315f..ad8ad9cc7f 100644
--- a/tensorflow_model_analysis/api/impl/evaluate.py
+++ b/tensorflow_model_analysis/api/impl/evaluate.py
@@ -18,17 +18,18 @@
 
 from __future__ import print_function
 
-import datetime
+
 
 
 import apache_beam as beam
 
 from tensorflow_model_analysis import constants
 from tensorflow_model_analysis import types
+from tensorflow_model_analysis.eval_saved_model import dofn
 from tensorflow_model_analysis.eval_saved_model import load
-from tensorflow_model_analysis.eval_saved_model import util
 from tensorflow_model_analysis.eval_saved_model.post_export_metrics import metric_keys
 from tensorflow_model_analysis.extractors import feature_extractor
+from tensorflow_model_analysis.extractors import predict_extractor
 from tensorflow_model_analysis.slicer import slicer
 from tensorflow_transform.beam import shared
 from tensorflow_model_analysis.types_compat import Any, Callable, Dict, Generator, List, Optional, Tuple
@@ -43,119 +44,6 @@
 _METRICS_NAMESPACE = 'tensorflow_model_analysis'
 
 
-def _make_construct_fn(  # pylint: disable=invalid-name
-    eval_saved_model_path,
-    add_metrics_callbacks,
-    model_load_seconds_distribution):
-  """Returns construct function for Shared for constructing EvalSavedModel."""
-
-  def construct():  # pylint: disable=invalid-name
-    """Function for constructing a EvalSavedModel."""
-    start_time = datetime.datetime.now()
-    result = load.EvalSavedModel(eval_saved_model_path)
-    if add_metrics_callbacks:
-      features_dict, predictions_dict, labels_dict = (
-          result.get_features_predictions_labels_dicts())
-      features_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
-          features_dict)
-      predictions_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
-          predictions_dict)
-      labels_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(labels_dict)
-      with result.graph_as_default():
-        metric_ops = {}
-        for add_metrics_callback in add_metrics_callbacks:
-          new_metric_ops = add_metrics_callback(features_dict, predictions_dict,
-                                                labels_dict)
-          overlap = set(new_metric_ops.keys()) & set(metric_ops.keys())
-          if overlap:
-            raise ValueError('metric keys should not conflict, but an '
-                             'earlier callback already added the metrics '
-                             'named %s' % overlap)
-          metric_ops.update(new_metric_ops)
-        result.register_additional_metric_ops(metric_ops)
-    end_time = datetime.datetime.now()
-    model_load_seconds_distribution.update(
-        int((end_time - start_time).total_seconds()))
-    return result
-
-  return construct
-
-
-class _EvalSavedModelDoFn(beam.DoFn):
-  """Abstract class for DoFns that load the EvalSavedModel and use it."""
-
-  def __init__(self, eval_saved_model_path,
-               add_metrics_callbacks,
-               shared_handle):
-    self._eval_saved_model_path = eval_saved_model_path
-    self._add_metrics_callbacks = add_metrics_callbacks
-    self._shared_handle = shared_handle
-    self._eval_saved_model = None  # type: load.EvalSavedModel
-    self._model_load_seconds = beam.metrics.Metrics.distribution(
-        _METRICS_NAMESPACE, 'model_load_seconds')
-
-  def start_bundle(self):
-    self._eval_saved_model = self._shared_handle.acquire(
-        _make_construct_fn(self._eval_saved_model_path,
-                           self._add_metrics_callbacks,
-                           self._model_load_seconds))
-
-  def process(self, elem):
-    raise NotImplementedError('not implemented')
-
-
-@beam.typehints.with_input_types(beam.typehints.List[types.ExampleAndExtracts])
-@beam.typehints.with_output_types(beam.typehints.Any)
-class _PredictionDoFn(_EvalSavedModelDoFn):
-  """A DoFn that loads the model and predicts."""
-
-  def __init__(self, eval_saved_model_path,
-               add_metrics_callbacks,
-               shared_handle):
-    super(_PredictionDoFn, self).__init__(eval_saved_model_path,
-                                          add_metrics_callbacks, shared_handle)
-    self._predict_batch_size = beam.metrics.Metrics.distribution(
-        _METRICS_NAMESPACE, 'predict_batch_size')
-    self._num_instances = beam.metrics.Metrics.counter(_METRICS_NAMESPACE,
-                                                       'num_instances')
-
-  def process(self, element
-             ):
-    batch_size = len(element)
-    self._predict_batch_size.update(batch_size)
-    self._num_instances.inc(batch_size)
-    serialized_examples = [x.example for x in element]
-
-    # Compute FeaturesPredictionsLabels for each serialized_example
-    for example_and_extracts, fpl in zip(
-        element,
-        self._eval_saved_model.predict_list(serialized_examples)):
-      example_and_extracts.extracts[
-          constants.FEATURES_PREDICTIONS_LABELS_KEY] = fpl
-
-    return element
-
-
-@beam.typehints.with_input_types(bytes)
-@beam.typehints.with_output_types(beam.typehints.Any)
-@beam.ptransform_fn
-def _Predict(  # pylint: disable=invalid-name
-    examples,
-    eval_saved_model_path,
-    desired_batch_size = None):
-  batch_args = {}
-  if desired_batch_size:
-    batch_args = dict(
-        min_batch_size=desired_batch_size, max_batch_size=desired_batch_size)
-  return (examples
-          | 'Batch' >> beam.BatchElements(**batch_args)
-          | beam.ParDo(
-              _PredictionDoFn(
-                  eval_saved_model_path=eval_saved_model_path,
-                  add_metrics_callbacks=None,
-                  shared_handle=shared.Shared())))
-
-
 @beam.typehints.with_input_types(beam.typehints.Any)
 @beam.typehints.with_output_types(
     beam.typehints.Tuple[_BeamSliceKeyType, beam.typehints.Any])
@@ -181,12 +69,12 @@ def process(self, element):
     self._post_slice_num_instances.inc(slice_count)
 
 
+@beam.ptransform_fn
 @beam.typehints.with_input_types(
     beam.typehints.Tuple[types.DictOfTensorType, beam.typehints.Any]
 )
 @beam.typehints.with_output_types(
-    beam.typehints.Tuple[_BeamSliceKeyType, beam.typehints.Any])
-@beam.ptransform_fn  # pylint: disable=invalid-name
+    beam.typehints.Tuple[_BeamSliceKeyType, beam.typehints.Any])  # pylint: disable=invalid-name
 def _Slice(intro_result,
            slice_spec):
   return intro_result | beam.ParDo(_SliceDoFn(slice_spec))
@@ -296,9 +184,9 @@ def _start_bundle(self):
     # There's no initialisation method for CombineFns.
     # See BEAM-3736: Add SetUp() and TearDown() for CombineFns.
     self._eval_saved_model = self._shared_handle.acquire(
-        _make_construct_fn(self._eval_saved_model_path,
-                           self._add_metrics_callbacks,
-                           self._model_load_seconds))
+        dofn.make_construct_fn(self._eval_saved_model_path,
+                               self._add_metrics_callbacks,
+                               self._model_load_seconds))
 
   def _maybe_do_batch(self, accumulator,
                       force = False):
@@ -357,11 +245,11 @@ def extract_output(self,
     return accumulator.metric_variables
 
 
+@beam.ptransform_fn
 @beam.typehints.with_input_types(
     beam.typehints.Tuple[_BeamSliceKeyType, beam.typehints.Any])
 @beam.typehints.with_output_types(beam.typehints.Tuple[
     _BeamSliceKeyType, beam.typehints.List[beam.typehints.Any]])
-@beam.ptransform_fn
 def _Aggregate(  # pylint: disable=invalid-name
     slice_result,
     eval_saved_model_path,
@@ -383,7 +271,7 @@ def _Aggregate(  # pylint: disable=invalid-name
     _BeamSliceKeyType, beam.typehints.List[beam.typehints.Any]])
 # No typehint for output type, since it's a multi-output DoFn result that
 # Beam doesn't support typehints for yet (BEAM-3280).
-class _ExtractOutputDoFn(_EvalSavedModelDoFn):
+class _ExtractOutputDoFn(dofn.EvalSavedModelDoFn):
   """A DoFn that extracts the metrics output."""
 
   OUTPUT_TAG_METRICS = 'tag_metrics'
@@ -404,14 +292,14 @@ def process(self, element
 
     yield (slice_key, slicing_metrics)
     if plots:
-      yield beam.pvalue.TaggedOutput(self.OUTPUT_TAG_PLOTS, (slice_key, plots))
+      yield beam.pvalue.TaggedOutput(self.OUTPUT_TAG_PLOTS, (slice_key, plots))  # pytype: disable=bad-return-type
 
 
+@beam.ptransform_fn
 @beam.typehints.with_input_types(beam.typehints.Tuple[
     _BeamSliceKeyType, beam.typehints.List[beam.typehints.Any]])
 # No typehint for output type, since it's a multi-output DoFn result that
 # Beam doesn't support typehints for yet (BEAM-3280).
-@beam.ptransform_fn
 def _ExtractOutput(  # pylint: disable=invalid-name
     aggregate_result, eval_saved_model_path,
     add_metrics_callbacks
@@ -485,45 +373,46 @@ def add_metrics_callback(features_dict, predictions_dict, labels):
     slice_spec = [slicer.SingleSliceSpec()]
 
   # pylint: disable=no-value-for-parameter
-  return (examples
-          # Our diagnostic outputs, pass types.ExampleAndExtracts throughout,
-          # however our aggregating functions do not use this interface.
-          | beam.Map(lambda x: types.ExampleAndExtracts(example=x, extracts={}))
-
-          # Map function which loads and runs the eval_saved_model against every
-          # example, yielding an types.ExampleAndExtracts containing a
-          # FeaturesPredictionsLabels value (where key is 'fpl').
-          | 'Predict' >> _Predict(
-              eval_saved_model_path=eval_saved_model_path,
-              desired_batch_size=desired_batch_size)
-
-          # Unwrap the types.ExampleAndExtracts.
-          # The rest of this pipeline expects FeaturesPredictionsLabels
-          | beam.Map(lambda x:  # pylint: disable=g-long-lambda
-                     x.extracts[constants.FEATURES_PREDICTIONS_LABELS_KEY])
-
-          # Input: one example fpl at a time
-          # Output: one fpl example per slice key (notice that the example turns
-          #         into n, replicated once per applicable slice key)
-          | 'Slice' >> _Slice(slice_spec)
-
-          # Each slice key lands on one shard where metrics are computed for all
-          # examples in that shard -- the "map" and "reduce" parts of the
-          # computation happen within this shard.
-          # Output: Tuple[slicer.SliceKeyType, MetricVariablesType]
-          | 'Aggregate' >> _Aggregate(
-              eval_saved_model_path=eval_saved_model_path,
-              add_metrics_callbacks=add_metrics_callbacks,
-              desired_batch_size=desired_batch_size)
+  return (
+      examples
+      # Our diagnostic outputs, pass types.ExampleAndExtracts throughout,
+      # however our aggregating functions do not use this interface.
+      | beam.Map(lambda x: types.ExampleAndExtracts(example=x, extracts={}))
+
+      # Map function which loads and runs the eval_saved_model against every
+      # example, yielding an types.ExampleAndExtracts containing a
+      # FeaturesPredictionsLabels value (where key is 'fpl').
+      | 'Predict' >> predict_extractor.TFMAPredict(
+          eval_saved_model_path=eval_saved_model_path,
+          desired_batch_size=desired_batch_size)
+
+      # Unwrap the types.ExampleAndExtracts.
+      # The rest of this pipeline expects FeaturesPredictionsLabels
+      | beam.Map(lambda x:  # pylint: disable=g-long-lambda
+                 x.extracts[constants.FEATURES_PREDICTIONS_LABELS_KEY])
+
+      # Input: one example fpl at a time
+      # Output: one fpl example per slice key (notice that the example turns
+      #         into n, replicated once per applicable slice key)
+      | 'Slice' >> _Slice(slice_spec)
+
+      # Each slice key lands on one shard where metrics are computed for all
+      # examples in that shard -- the "map" and "reduce" parts of the
+      # computation happen within this shard.
+      # Output: Tuple[slicer.SliceKeyType, MetricVariablesType]
+      | 'Aggregate' >> _Aggregate(
+          eval_saved_model_path=eval_saved_model_path,
+          add_metrics_callbacks=add_metrics_callbacks,
+          desired_batch_size=desired_batch_size)
 
-          # Different metrics for a given slice key are brought together.
-          | 'ExtractOutput' >> _ExtractOutput(eval_saved_model_path,
-                                              add_metrics_callbacks))
+      # Different metrics for a given slice key are brought together.
+      | 'ExtractOutput' >> _ExtractOutput(eval_saved_model_path,
+                                          add_metrics_callbacks))
 
 
+@beam.ptransform_fn
 @beam.typehints.with_input_types(bytes)
 @beam.typehints.with_output_types(beam.typehints.Any)
-@beam.ptransform_fn
 def BuildDiagnosticTable(
     # pylint: disable=invalid-name
     examples,
@@ -543,7 +432,8 @@ def BuildDiagnosticTable(
     PCollection of ExampleAndExtracts
   """
   return (examples
-          | 'ToExampleAndExtracts' >> beam.Map(
-              lambda x: types.ExampleAndExtracts(example=x, extracts={}))
-          | 'Predict' >> _Predict(eval_saved_model_path, desired_batch_size)
+          | 'ToExampleAndExtracts' >>
+          beam.Map(lambda x: types.ExampleAndExtracts(example=x, extracts={}))
+          | 'Predict' >> predict_extractor.TFMAPredict(eval_saved_model_path,
+                                                       desired_batch_size)
           | 'ExtractFeatures' >> feature_extractor.ExtractFeatures())
diff --git a/tensorflow_model_analysis/api/model_eval_lib.py b/tensorflow_model_analysis/api/model_eval_lib.py
index bf8a8b20d3..e8aef8db61 100644
--- a/tensorflow_model_analysis/api/model_eval_lib.py
+++ b/tensorflow_model_analysis/api/model_eval_lib.py
@@ -123,9 +123,9 @@ def load_eval_result(output_path):
 def _assert_tensorflow_version():
   # Fail with a clear error in case we are not using a compatible TF version.
   major, minor, _ = tf.__version__.split('.')
-  if int(major) != 1 or int(minor) < 9:
+  if int(major) != 1 or int(minor) < 10:
     raise RuntimeError(
-        'Tensorflow version >= 1.9, < 2 is required. Found (%s). Please '
+        'Tensorflow version >= 1.10, < 2 is required. Found (%s). Please '
         'install the latest 1.x version from '
         'https://github.com/tensorflow/tensorflow. ' % tf.__version__)
 
diff --git a/tensorflow_model_analysis/contrib/BUILD b/tensorflow_model_analysis/contrib/BUILD
deleted file mode 100644
index ced633f920..0000000000
--- a/tensorflow_model_analysis/contrib/BUILD
+++ /dev/null
@@ -1,31 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-exports_files(["LICENSE"])
-
-py_library(
-    name = "model_eval_lib",
-    srcs = ["model_eval_lib.py"],
-    deps = [
-        "//tensorflow_model_analysis:constants",
-        "//tensorflow_model_analysis:expect_apache_beam_installed",  # b/73825929
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/api/impl:api_types",
-        "//tensorflow_model_analysis/api/impl:evaluate",
-    ],
-)
-
-py_test(
-    name = "model_eval_lib_test",
-    srcs = ["model_eval_lib_test.py"],
-    deps = [
-        ":model_eval_lib",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:testutil",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:csv_linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:custom_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics",
-        "//tensorflow_model_analysis/slicer",
-    ],
-)
diff --git a/tensorflow_model_analysis/contrib/model_eval_lib.py b/tensorflow_model_analysis/contrib/model_eval_lib.py
index f767fba4b0..dcf6d86ed2 100644
--- a/tensorflow_model_analysis/contrib/model_eval_lib.py
+++ b/tensorflow_model_analysis/contrib/model_eval_lib.py
@@ -25,9 +25,9 @@
 from tensorflow_model_analysis.types_compat import Optional
 
 
+@beam.ptransform_fn
 @beam.typehints.with_input_types(bytes)
 @beam.typehints.with_output_types(beam.typehints.Any)
-@beam.ptransform_fn
 def BuildDiagnosticTable(  # pylint: disable=invalid-name
     examples,
     eval_saved_model_path,
diff --git a/tensorflow_model_analysis/eval_saved_model/BUILD b/tensorflow_model_analysis/eval_saved_model/BUILD
deleted file mode 100644
index 4135cdf5bf..0000000000
--- a/tensorflow_model_analysis/eval_saved_model/BUILD
+++ /dev/null
@@ -1,141 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-py_library(
-    name = "constants",
-    srcs = ["constants.py"],
-    deps = [],
-)
-
-py_library(
-    name = "export",
-    srcs = ["export.py"],
-    deps = [
-        ":constants",
-        ":encoding",
-        ":util",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:types",
-        "//tensorflow_model_analysis:version",
-    ],
-)
-
-py_library(
-    name = "exporter",
-    srcs = ["exporter.py"],
-    deps = [
-        ":export",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-    ],
-)
-
-py_library(
-    name = "encoding",
-    srcs = ["encoding.py"],
-    deps = [
-        "//google/protobuf:any_py_pb2",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:types",
-        "//third_party/tensorflow/core:protos_all_py_pb2",
-    ],
-)
-
-py_test(
-    name = "encoding_test",
-    srcs = ["encoding_test.py"],
-    deps = [
-        ":encoding",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-    ],
-)
-
-py_library(
-    name = "load",
-    srcs = ["load.py"],
-    deps = [
-        ":constants",
-        ":encoding",
-        ":graph_ref",
-        ":util",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:types",
-        "//tensorflow_model_analysis:util",
-    ],
-)
-
-py_library(
-    name = "graph_ref",
-    srcs = ["graph_ref.py"],
-    deps = [
-        ":encoding",
-        "//google/protobuf:any_py_pb2",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:types",
-        "//third_party/tensorflow/core:protos_all_py_pb2",
-    ],
-)
-
-py_test(
-    name = "graph_ref_test",
-    srcs = ["graph_ref_test.py"],
-    deps = [
-        ":graph_ref",
-        "//third_party/tensorflow/core:protos_all_py_pb2",
-    ],
-)
-
-py_test(
-    name = "integration_test",
-    srcs = ["integration_test.py"],
-    shard_count = 6,
-    deps = [
-        ":encoding",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:load",
-        "//tensorflow_model_analysis/eval_saved_model:testutil",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:csv_linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:custom_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:dnn_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fake_sequence_to_prediction",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:linear_classifier_multivalent",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:multi_head",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics:metrics",
-        "//third_party/tensorflow/core:protos_all_py_pb2",
-    ],
-)
-
-py_library(
-    name = "testutil",
-    testonly = 1,
-    deps = [":testutil_non_testonly"],
-)
-
-py_library(
-    name = "testutil_non_testonly",
-    srcs = ["testutil.py"],
-    deps = [
-        ":util",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-    ],
-)
-
-py_library(
-    name = "util",
-    srcs = ["util.py"],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:types",
-        "//third_party/py/numpy",
-    ],
-)
-
-py_test(
-    name = "util_test",
-    srcs = ["util_test.py"],
-    deps = [
-        ":testutil",
-        ":util",
-        "//third_party/tensorflow/core:protos_all_py_pb2",
-    ],
-)
diff --git a/tensorflow_model_analysis/eval_saved_model/dofn.py b/tensorflow_model_analysis/eval_saved_model/dofn.py
new file mode 100644
index 0000000000..fc947d59b9
--- /dev/null
+++ b/tensorflow_model_analysis/eval_saved_model/dofn.py
@@ -0,0 +1,90 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""DoFns that load EvalSavedModel and use it."""
+
+
+
+import datetime
+
+import apache_beam as beam
+from tensorflow_model_analysis import types
+from tensorflow_model_analysis.eval_saved_model import load
+from tensorflow_model_analysis.eval_saved_model import util
+from tensorflow_transform.beam import shared
+
+from tensorflow_model_analysis.types_compat import Callable, Dict, List, Optional, Tuple
+
+_METRICS_NAMESPACE = 'tensorflow_model_analysis'
+
+
+
+def make_construct_fn(  # pylint: disable=invalid-name
+    eval_saved_model_path,
+    add_metrics_callbacks,
+    model_load_seconds_distribution):
+  """Returns construct function for Shared for constructing EvalSavedModel."""
+
+  def construct():  # pylint: disable=invalid-name
+    """Function for constructing a EvalSavedModel."""
+    start_time = datetime.datetime.now()
+    result = load.EvalSavedModel(eval_saved_model_path)
+    if add_metrics_callbacks:
+      features_dict, predictions_dict, labels_dict = (
+          result.get_features_predictions_labels_dicts())
+      features_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
+          features_dict)
+      predictions_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(
+          predictions_dict)
+      labels_dict = util.wrap_tensor_or_dict_of_tensors_in_identity(labels_dict)
+      with result.graph_as_default():
+        metric_ops = {}
+        for add_metrics_callback in add_metrics_callbacks:
+          new_metric_ops = add_metrics_callback(features_dict, predictions_dict,
+                                                labels_dict)
+          overlap = set(new_metric_ops.keys()) & set(metric_ops.keys())
+          if overlap:
+            raise ValueError('metric keys should not conflict, but an '
+                             'earlier callback already added the metrics '
+                             'named %s' % overlap)
+          metric_ops.update(new_metric_ops)
+        result.register_additional_metric_ops(metric_ops)
+    end_time = datetime.datetime.now()
+    model_load_seconds_distribution.update(
+        int((end_time - start_time).total_seconds()))
+    return result
+
+  return construct
+
+
+class EvalSavedModelDoFn(beam.DoFn):
+  """Abstract class for DoFns that load the EvalSavedModel and use it."""
+
+  def __init__(self, eval_saved_model_path,
+               add_metrics_callbacks,
+               shared_handle):
+    self._eval_saved_model_path = eval_saved_model_path
+    self._add_metrics_callbacks = add_metrics_callbacks
+    self._shared_handle = shared_handle
+    self._eval_saved_model = None  # type: load.EvalSavedModel
+    self._model_load_seconds = beam.metrics.Metrics.distribution(
+        _METRICS_NAMESPACE, 'model_load_seconds')
+
+  def start_bundle(self):
+    self._eval_saved_model = self._shared_handle.acquire(
+        make_construct_fn(self._eval_saved_model_path,
+                          self._add_metrics_callbacks,
+                          self._model_load_seconds))
+
+  def process(self, elem):
+    raise NotImplementedError('not implemented')
diff --git a/tensorflow_model_analysis/eval_saved_model/example_trainers/BUILD b/tensorflow_model_analysis/eval_saved_model/example_trainers/BUILD
deleted file mode 100644
index 1f2abd33c3..0000000000
--- a/tensorflow_model_analysis/eval_saved_model/example_trainers/BUILD
+++ /dev/null
@@ -1,169 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-py_library(
-    name = "custom_estimator",
-    srcs = [
-        "custom_estimator.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "fixed_prediction_estimator",
-    srcs = [
-        "fixed_prediction_estimator.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "fixed_prediction_classifier",
-    srcs = [
-        "fixed_prediction_classifier.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "fixed_prediction_estimator_extra_fields",
-    srcs = [
-        "fixed_prediction_estimator_extra_fields.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "fixed_prediction_classifier_extra_fields",
-    srcs = [
-        "fixed_prediction_classifier_extra_fields.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "linear_classifier",
-    srcs = [
-        "linear_classifier.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "dnn_classifier",
-    srcs = [
-        "dnn_classifier.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "linear_classifier_multivalent",
-    srcs = [
-        "linear_classifier_multivalent.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "linear_regressor",
-    srcs = [
-        "linear_regressor.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "dnn_regressor",
-    srcs = [
-        "dnn_regressor.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "fake_sequence_to_prediction",
-    srcs = [
-        "fake_sequence_to_prediction.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model:util",
-    ],
-)
-
-py_library(
-    name = "util",
-    srcs = [
-        "util.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model:util",
-    ],
-)
-
-py_library(
-    name = "csv_linear_classifier",
-    srcs = [
-        "csv_linear_classifier.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
-
-py_library(
-    name = "multi_head",
-    srcs = [
-        "multi_head.py",
-    ],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:export",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:util",
-    ],
-)
diff --git a/tensorflow_model_analysis/eval_saved_model/example_trainers/fixed_prediction_estimator.py b/tensorflow_model_analysis/eval_saved_model/example_trainers/fixed_prediction_estimator.py
index 793f518135..203919bab1 100644
--- a/tensorflow_model_analysis/eval_saved_model/example_trainers/fixed_prediction_estimator.py
+++ b/tensorflow_model_analysis/eval_saved_model/example_trainers/fixed_prediction_estimator.py
@@ -30,16 +30,25 @@
 from tensorflow.python.estimator.canned import prediction_keys
 
 
-def simple_fixed_prediction_estimator(export_path, eval_export_path):
+def simple_fixed_prediction_estimator(
+    export_path,
+    eval_export_path,
+    output_prediction_key=prediction_keys.PredictionKeys.PREDICTIONS):
   """Exports a simple fixed prediction estimator."""
 
   def model_fn(features, labels, mode, params):
     """Model function for custom estimator."""
     del params
     predictions = features['prediction']
-    predictions_dict = {
-        prediction_keys.PredictionKeys.PREDICTIONS: predictions,
-    }
+
+    if output_prediction_key is not None:
+      predictions_dict = {
+          output_prediction_key: predictions,
+      }
+    else:
+      # For simulating Estimators which don't return a predictions dict in
+      # EVAL mode.
+      predictions_dict = {}
 
     if mode == tf.estimator.ModeKeys.PREDICT:
       return tf.estimator.EstimatorSpec(
diff --git a/tensorflow_model_analysis/eval_saved_model/post_export_metrics/BUILD b/tensorflow_model_analysis/eval_saved_model/post_export_metrics/BUILD
deleted file mode 100644
index c82730c7ab..0000000000
--- a/tensorflow_model_analysis/eval_saved_model/post_export_metrics/BUILD
+++ /dev/null
@@ -1,50 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-py_library(
-    name = "post_export_metrics",
-    srcs = ["post_export_metrics.py"],
-    deps = [
-        ":metric_keys",
-        ":metrics",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:types",
-        "//tensorflow_model_analysis/eval_saved_model:encoding",
-        "//tensorflow_model_analysis/proto:metrics_for_slice_py_pb2",
-    ],
-)
-
-py_library(
-    name = "metric_keys",
-    srcs = ["metric_keys.py"],
-)
-
-py_library(
-    name = "metrics",
-    srcs = ["metrics.py"],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:types",
-    ],
-)
-
-py_test(
-    name = "post_export_metrics_test",
-    srcs = ["post_export_metrics_test.py"],
-    shard_count = 5,
-    deps = [
-        ":post_export_metrics",
-        "//tensorflow_model_analysis:expect_apache_beam_installed",  # b/73825929
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/api/impl:evaluate",
-        "//tensorflow_model_analysis/eval_saved_model:testutil",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:dnn_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:dnn_regressor",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_classifier_extra_fields",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_estimator_extra_fields",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:linear_regressor",
-        "//tensorflow_model_analysis/proto:metrics_for_slice_py_pb2",
-    ],
-)
diff --git a/tensorflow_model_analysis/eval_saved_model/post_export_metrics/post_export_metrics.py b/tensorflow_model_analysis/eval_saved_model/post_export_metrics/post_export_metrics.py
index cd857a871f..081304a4e8 100644
--- a/tensorflow_model_analysis/eval_saved_model/post_export_metrics/post_export_metrics.py
+++ b/tensorflow_model_analysis/eval_saved_model/post_export_metrics/post_export_metrics.py
@@ -110,10 +110,8 @@ def _get_prediction_tensor(
     predictions_dict: Predictions dictionary.
 
   Returns:
-    Predictions tensor.
-
-  Raises:
-    KeyError: No expected keys are found in predictions_dict.
+    Predictions tensor, or None if none of the expected keys are found in
+    the predictions_dict.
   """
   if types.is_tensor(predictions_dict):
     return predictions_dict
@@ -126,8 +124,8 @@ def _get_prediction_tensor(
     ref_tensor = predictions_dict.get(key)
     if ref_tensor is not None:
       return ref_tensor
-  raise KeyError('cannot find any keys %s in predictions_dict %s.' %
-                 (key_precedence, predictions_dict))
+
+  return None
 
 
 def _check_labels(labels_dict):
@@ -138,7 +136,9 @@ def _check_labels(labels_dict):
 
 def _check_predictions(predictions_dict):
   """Raise KeyError if the predictions cannot be understood."""
-  _get_prediction_tensor(predictions_dict)
+  if _get_prediction_tensor(predictions_dict) is None:
+    raise KeyError('cannot find any of the standard keysin predictions_dict %s.'
+                   % (predictions_dict))
 
 
 def _check_labels_and_predictions(
@@ -274,19 +274,62 @@ def populate_plots_and_pop(
 
 @_export('example_count')
 class _ExampleCount(_PostExportMetric):
-  """Metric that counts the number of examples processed."""
+  """Metric that counts the number of examples processed.
+
+  We get the example count by looking at the predictions dictionary and picking
+  a reference Tensor. If we can find a standard key (e.g.
+  PredictionKeys.LOGISTIC, etc), we use that as the reference Tensor. Otherwise,
+  we just use the first key in sorted order from one of the dictionaries
+  (predictions, labels) as the reference Tensor.
+
+  We assume the first dimension is the batch size, and take that to be the
+  number of examples in the batch.
+  """
 
   def check_compatibility(self, features_dict,
                           predictions_dict,
                           labels_dict):
-    _check_predictions(predictions_dict)
+    pass
 
   def get_metric_ops(self, features_dict,
                      predictions_dict,
                      labels_dict
                     ):
     ref_tensor = _get_prediction_tensor(predictions_dict)
-    return {metric_keys.EXAMPLE_COUNT: tf.contrib.metrics.count(ref_tensor)}
+    if ref_tensor is None:
+      # Note that if predictions_dict is a Tensor and not a dict,
+      # get_predictions_tensor will return predictions_dict, so if we get
+      # here, if means that predictions_dict is a dict without any of the
+      # standard keys.
+      #
+      # If we can't get any of standard keys, then pick the first key
+      # in alphabetical order if the predictions dict is non-empty.
+      # If the predictions dict is empty, try the labels dict.
+      # If that is empty too, default to the empty Tensor.
+      tf.logging.info(
+          'ExampleCount post export metric: could not find any of '
+          'the standard keys in predictions_dict (keys were: %s)',
+          predictions_dict.keys())
+      if predictions_dict is not None and predictions_dict.keys():
+        first_key = sorted(predictions_dict.keys())[0]
+        ref_tensor = predictions_dict[first_key]
+        tf.logging.info('Using the first key from predictions_dict: %s',
+                        first_key)
+      elif labels_dict is not None:
+        if types.is_tensor(labels_dict):
+          ref_tensor = labels_dict
+          tf.logging.info('Using the labels Tensor')
+        elif labels_dict.keys():
+          first_key = sorted(labels_dict.keys())[0]
+          ref_tensor = labels_dict[first_key]
+          tf.logging.info('Using the first key from labels_dict: %s', first_key)
+
+      if ref_tensor is None:
+        tf.logging.info('Could not find a reference Tensor for example count. '
+                        'Defaulting to the empty Tensor.')
+        ref_tensor = tf.constant([])
+
+    return {metric_keys.EXAMPLE_COUNT: metrics.total(tf.shape(ref_tensor)[0])}
 
 
 @_export('example_weight')
diff --git a/tensorflow_model_analysis/eval_saved_model/post_export_metrics/post_export_metrics_test.py b/tensorflow_model_analysis/eval_saved_model/post_export_metrics/post_export_metrics_test.py
index 937bfbda0d..063727b012 100644
--- a/tensorflow_model_analysis/eval_saved_model/post_export_metrics/post_export_metrics_test.py
+++ b/tensorflow_model_analysis/eval_saved_model/post_export_metrics/post_export_metrics_test.py
@@ -82,6 +82,46 @@ def check_result(got):  # pylint: disable=invalid-name
     self._runTestWithCustomCheck(
         examples, eval_export_dir, metrics, custom_metrics_check=check_result)
 
+  def testExampleCountNoStandardKeys(self):
+    # Test ExampleCount with a custom Estimator that doesn't have any of the
+    # standard PredictionKeys.
+    temp_eval_export_dir = self._getEvalExportDir()
+    _, eval_export_dir = (
+        fixed_prediction_estimator.simple_fixed_prediction_estimator(
+            None, temp_eval_export_dir, output_prediction_key='non_standard'))
+    examples = [
+        self._makeExample(prediction=5.0, label=5.0),
+        self._makeExample(prediction=6.0, label=6.0),
+        self._makeExample(prediction=7.0, label=7.0),
+    ]
+    expected_values_dict = {
+        metric_keys.EXAMPLE_COUNT: 3.0,
+    }
+    self._runTest(examples, eval_export_dir, [
+        post_export_metrics.example_count(),
+    ], expected_values_dict)
+
+  def testExampleCountEmptyPredictionsDict(self):
+    # Test ExampleCount with a custom Estimator that has empty predictions dict.
+    # This is possible if the Estimator doesn't return the predictions dict
+    # in EVAL mode, but computes predictions and feeds them into the metrics
+    # internally.
+    temp_eval_export_dir = self._getEvalExportDir()
+    _, eval_export_dir = (
+        fixed_prediction_estimator.simple_fixed_prediction_estimator(
+            None, temp_eval_export_dir, output_prediction_key=None))
+    examples = [
+        self._makeExample(prediction=5.0, label=5.0),
+        self._makeExample(prediction=6.0, label=6.0),
+        self._makeExample(prediction=7.0, label=7.0),
+    ]
+    expected_values_dict = {
+        metric_keys.EXAMPLE_COUNT: 3.0,
+    }
+    self._runTest(examples, eval_export_dir, [
+        post_export_metrics.example_count(),
+    ], expected_values_dict)
+
   def testPostExportMetricsLinearClassifier(self):
     temp_eval_export_dir = self._getEvalExportDir()
     _, eval_export_dir = linear_classifier.simple_linear_classifier(
diff --git a/tensorflow_model_analysis/extractors/BUILD b/tensorflow_model_analysis/extractors/BUILD
deleted file mode 100644
index f54bf6d9f7..0000000000
--- a/tensorflow_model_analysis/extractors/BUILD
+++ /dev/null
@@ -1,29 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-exports_files(["LICENSE"])
-
-py_library(
-    name = "feature_extractor",
-    srcs = ["feature_extractor.py"],
-    deps = [
-        "//tensorflow_model_analysis:constants",
-        "//tensorflow_model_analysis:expect_apache_beam_installed",  # b/73825929
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/api/impl:api_types",
-        "//tensorflow_model_analysis/eval_saved_model:load",
-    ],
-)
-
-py_test(
-    name = "feature_extractor_test",
-    srcs = ["feature_extractor_test.py"],
-    deps = [
-        ":feature_extractor",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:testutil",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:csv_linear_classifier",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:custom_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:fixed_prediction_estimator",
-        "//tensorflow_model_analysis/eval_saved_model/example_trainers:linear_classifier",
-    ],
-)
diff --git a/tensorflow_model_analysis/extractors/feature_extractor.py b/tensorflow_model_analysis/extractors/feature_extractor.py
index 50c322755f..eebc0f0d41 100644
--- a/tensorflow_model_analysis/extractors/feature_extractor.py
+++ b/tensorflow_model_analysis/extractors/feature_extractor.py
@@ -37,16 +37,17 @@
 def _AugmentExtracts(
     fpl_dict,
     example_and_extracts):
-  """Augments The ExampleAndExtracts with FeaturesPredictionsLabels.
+  """Augments the ExampleAndExtracts with FeaturesPredictionsLabels.
 
   Args:
     fpl_dict: The dictionary returned by evaluate._Predict()
-    example_and_extracts: The ExampleAndExtracts to be augmented -- note that
-      this variable modified (ie both an input and output)
+    example_and_extracts: The ExampleAndExtracts to be augmented. This is
+      mutated in-place.
+
   Raises:
-    TypeError: if the FeaturesPredictionsLabels is corrupt.
+    TypeError: If the FeaturesPredictionsLabels is corrupt.
   """
-  for name, val in fpl_dict.iteritems():
+  for name, val in fpl_dict.items():
     val = val.get(encoding.NODE_SUFFIX)
 
     if isinstance(val, tf.SparseTensorValue):
@@ -62,7 +63,9 @@ def _AugmentExtracts(
           name=name, value=val)
 
     else:
-      raise TypeError('Unexpected fpl type: %s' % str(val))
+      raise TypeError(
+          'Dictionary item with key %s, value %s had unexpected type %s' %
+          (name, val, type(val)))
 
 
 def _MaterializeFeatures(
@@ -76,35 +79,37 @@ def _MaterializeFeatures(
     example_and_extracts: The ExampleAndExtracts to be augmented
 
   Returns:
-    Reference to augmented ExampleAndExtracts.
+    Returns an augmented ExampleAndExtracts (which is a shallow copy of
+    the original ExampleAndExtracts, so the original isn't mutated)
 
   Raises:
     RuntimeError: When _Predict() didn't populate the 'fpl' key.
   """
-  fpl = example_and_extracts.extracts.get(
-      constants.FEATURES_PREDICTIONS_LABELS_KEY)
+  # Make a a shallow copy, so we don't mutate the original.
+  result = example_and_extracts.create_copy_with_shallow_copy_of_extracts()
+
+  fpl = result.extracts.get(constants.FEATURES_PREDICTIONS_LABELS_KEY)
   if not fpl:
-    raise RuntimeError(
-        'fpl missing, Please ensure _Predict() was called.')
+    raise RuntimeError('FPL missing, Please ensure _Predict() was called.')
 
   if not isinstance(fpl, load.FeaturesPredictionsLabels):
-    raise RuntimeError(
-        'Expected FPL to be instance of FeaturesPredictionsLabel. FPL was: %s'
-        % str(fpl))
+    raise TypeError(
+        'Expected FPL to be instance of FeaturesPredictionsLabel. FPL was: %s '
+        'of type %s' % (str(fpl), type(fpl)))
 
   # We disable pytyping here because we know that 'fpl' key corresponds to a
   # non-materialized column.
   # pytype: disable=attribute-error
-  _AugmentExtracts(fpl.features, example_and_extracts)
-  _AugmentExtracts(fpl.predictions, example_and_extracts)
-  _AugmentExtracts(fpl.labels, example_and_extracts)
-  return example_and_extracts
+  _AugmentExtracts(fpl.features, result)
+  _AugmentExtracts(fpl.predictions, result)
+  _AugmentExtracts(fpl.labels, result)
+  return result
   # pytype: enable=attribute-error
 
 
+@beam.ptransform_fn
 @beam.typehints.with_input_types(beam.typehints.Any)
 @beam.typehints.with_output_types(beam.typehints.Any)
-@beam.ptransform_fn
 def ExtractFeatures(
     examples_and_extracts):
   """Builds MaterializedColumn extracts from FPL created in evaluate.Predict().
diff --git a/tensorflow_model_analysis/extractors/feature_extractor_test.py b/tensorflow_model_analysis/extractors/feature_extractor_test.py
index 6e4c2d1033..1da0152e77 100644
--- a/tensorflow_model_analysis/extractors/feature_extractor_test.py
+++ b/tensorflow_model_analysis/extractors/feature_extractor_test.py
@@ -57,7 +57,7 @@ def testMaterializeFeaturesBadFPL(self):
     example_and_extracts = types.ExampleAndExtracts(
         example=example1.SerializeToString(),
         extracts={'fpl': 123})
-    self.assertRaises(RuntimeError, feature_extractor._MaterializeFeatures,
+    self.assertRaises(TypeError, feature_extractor._MaterializeFeatures,
                       example_and_extracts)
 
   def testMaterializeFeaturesNoMaterializedColumns(self):
diff --git a/tensorflow_model_analysis/extractors/predict_extractor.py b/tensorflow_model_analysis/extractors/predict_extractor.py
new file mode 100644
index 0000000000..15368150b9
--- /dev/null
+++ b/tensorflow_model_analysis/extractors/predict_extractor.py
@@ -0,0 +1,96 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Public API for performing evaluations using the EvalSavedModel."""
+
+from __future__ import absolute_import
+from __future__ import division
+
+from __future__ import print_function
+
+
+
+import apache_beam as beam
+
+from tensorflow_model_analysis import constants
+from tensorflow_model_analysis import types
+from tensorflow_model_analysis.eval_saved_model import dofn
+from tensorflow_transform.beam import shared
+from tensorflow_model_analysis.types_compat import Any, Callable, Dict, List, Optional, Tuple
+
+MetricVariablesType = List[Any]  # pylint: disable=invalid-name
+
+# For use in Beam type annotations, because Beam's support for Python types
+# in Beam type annotations is not complete.
+_BeamSliceKeyType = beam.typehints.Tuple[  # pylint: disable=invalid-name
+    beam.typehints.Tuple[bytes, beam.typehints.Union[bytes, int, float]], Ellipsis]
+
+_METRICS_NAMESPACE = 'tensorflow_model_analysis'
+
+
+@beam.typehints.with_input_types(beam.typehints.List[types.ExampleAndExtracts])
+@beam.typehints.with_output_types(beam.typehints.Any)
+class _TFMAPredictionDoFn(dofn.EvalSavedModelDoFn):
+  """A DoFn that loads the model and predicts."""
+
+  def __init__(self, eval_saved_model_path,
+               add_metrics_callbacks,
+               shared_handle):
+    super(_TFMAPredictionDoFn, self).__init__(
+        eval_saved_model_path, add_metrics_callbacks, shared_handle)
+    self._predict_batch_size = beam.metrics.Metrics.distribution(
+        _METRICS_NAMESPACE, 'predict_batch_size')
+    self._num_instances = beam.metrics.Metrics.counter(_METRICS_NAMESPACE,
+                                                       'num_instances')
+
+  def process(self, element
+             ):
+    result = []
+    batch_size = len(element)
+    self._predict_batch_size.update(batch_size)
+    self._num_instances.inc(batch_size)
+    serialized_examples = [x.example for x in element]
+
+    # Compute FeaturesPredictionsLabels for each serialized_example
+    for example_and_extracts, fpl in zip(
+        element, self._eval_saved_model.predict_list(serialized_examples)):
+
+      # Make a a shallow copy, so we don't mutate the original.
+      element_copy = (
+          example_and_extracts.create_copy_with_shallow_copy_of_extracts())
+      element_copy.extracts[constants.FEATURES_PREDICTIONS_LABELS_KEY] = fpl
+
+      result.append(element_copy)
+
+    return result
+
+
+@beam.ptransform_fn
+@beam.typehints.with_input_types(beam.typehints.List[types.ExampleAndExtracts])
+@beam.typehints.with_output_types(beam.typehints.Any)
+def TFMAPredict(  # pylint: disable=invalid-name
+    examples,
+    eval_saved_model_path,
+    desired_batch_size = None):
+  """A PTransform that adds predictions to ExamplesAndExtracts."""
+  batch_args = {}
+  if desired_batch_size:
+    batch_args = dict(
+        min_batch_size=desired_batch_size, max_batch_size=desired_batch_size)
+  return (examples
+          | 'Batch' >> beam.BatchElements(**batch_args)
+          | beam.ParDo(
+              _TFMAPredictionDoFn(
+                  eval_saved_model_path=eval_saved_model_path,
+                  add_metrics_callbacks=None,
+                  shared_handle=shared.Shared())))
diff --git a/tensorflow_model_analysis/extractors/predict_extractor_test.py b/tensorflow_model_analysis/extractors/predict_extractor_test.py
new file mode 100644
index 0000000000..44ab67c82d
--- /dev/null
+++ b/tensorflow_model_analysis/extractors/predict_extractor_test.py
@@ -0,0 +1,89 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Test for using the Evaluate API.
+
+Note that we actually train and export models within these tests.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+
+import apache_beam as beam
+
+from apache_beam.testing import util
+
+import tensorflow as tf
+
+from tensorflow_model_analysis import types
+from tensorflow_model_analysis.eval_saved_model import testutil
+from tensorflow_model_analysis.eval_saved_model.example_trainers import linear_classifier
+from tensorflow_model_analysis.extractors import predict_extractor
+
+
+class PredictExtractorTest(testutil.TensorflowModelAnalysisTest):
+
+  def _getEvalExportDir(self):
+    return os.path.join(self._getTempDir(), 'eval_export_dir')
+
+  def testPredict(self):
+    temp_eval_export_dir = self._getEvalExportDir()
+    _, eval_export_dir = linear_classifier.simple_linear_classifier(
+        None, temp_eval_export_dir)
+
+    with beam.Pipeline() as pipeline:
+      example1 = self._makeExample(age=3.0, language='english', label=1.0)
+      example2 = self._makeExample(age=3.0, language='chinese', label=0.0)
+      example3 = self._makeExample(age=4.0, language='english', label=1.0)
+      example4 = self._makeExample(age=5.0, language='chinese', label=0.0)
+
+      predict_extracts = (
+          pipeline
+          | beam.Create([
+              example1.SerializeToString(),
+              example2.SerializeToString(),
+              example3.SerializeToString(),
+              example4.SerializeToString()
+          ])
+          # Our diagnostic outputs, pass types.ExampleAndExtracts throughout,
+          # however our aggregating functions do not use this interface.
+          | beam.Map(lambda x: types.ExampleAndExtracts(example=x, extracts={}))
+          | 'Predict' >> predict_extractor.TFMAPredict(
+              eval_saved_model_path=eval_export_dir, desired_batch_size=3))
+
+      def check_result(got):
+        try:
+          self.assertEqual(4, len(got), 'got: %s' % got)
+          for item in got:
+            extracts_dict = item.extracts
+            self.assertTrue(extracts_dict.has_key('fpl'))
+            fpl = extracts_dict['fpl']
+            # Verify fpl contains features, probabilities, and correct labels.
+            self.assertIn('language', fpl.features)
+            self.assertIn('age', fpl.features)
+            self.assertIn('label', fpl.features)
+            self.assertIn('probabilities', fpl.predictions)
+            self.assertAlmostEqual(fpl.features['label'],
+                                   fpl.labels['__labels'])
+        except AssertionError as err:
+          raise util.BeamAssertException(err)
+
+      util.assert_that(predict_extracts, check_result)
+
+
+if __name__ == '__main__':
+  tf.test.main()
diff --git a/tensorflow_model_analysis/frontend/test/tfma-bounded-value_test.html b/tensorflow_model_analysis/frontend/test/tfma-bounded-value_test.html
index 6189ff399e..feaf25d88e 100644
--- a/tensorflow_model_analysis/frontend/test/tfma-bounded-value_test.html
+++ b/tensorflow_model_analysis/frontend/test/tfma-bounded-value_test.html
@@ -44,15 +44,36 @@
   test('setData', () => {
     element = fixture('plain-fixture');
     element.data = '{"lowerBound": 1, "upperBound": 2, "value": 1.5}';
-    assert.equal(flatten(element.textContent), '1.50000 (1.00000, 2.00000)');
+    checkText(element, '1.50000 (1.00000, 2.00000)');
   });
 
   test('createBoundedValueWithAttributes', () => {
     element = fixture('attributes-inlined-fixture');
-    assert.equal(flatten(element.textContent), '1.50000 (1.00000, 2.00000)');
+    debugger;
+    checkText(element, '1.50000 (1.00000, 2.00000)');
   });
 });
 
+/**
+ * Checks the given element contains the provided text.
+ * @param {!Element} element
+ * @param {string} expectedText
+ */
+function checkText(element, expectedText) {
+  // Use indexOf instead of full string match since Polymer 2.x will also contain the css inlined in
+  // the template.
+  assert.isTrue(flatten(getTextContent(element)).indexOf(expectedText) >= 0);
+}
+
+/**
+ * Extracts the text content of the element. For polymer 2.X, get the text from the shadow root.
+ * @param {!Element} element
+ * @return {string}
+ */
+function getTextContent(element) {
+  return element.textContent.trim() || element.root.textContent.trim();
+}
+
 /**
  * For string comparison purpose, "flattens" the given html by removing all
  * unnecessary newline and white space.
diff --git a/tensorflow_model_analysis/frontend/test/tfma-plot_test.html b/tensorflow_model_analysis/frontend/test/tfma-plot_test.html
index 206fff7103..bb93a21880 100644
--- a/tensorflow_model_analysis/frontend/test/tfma-plot_test.html
+++ b/tensorflow_model_analysis/frontend/test/tfma-plot_test.html
@@ -216,7 +216,7 @@
     ];
     const macroPrecisionRecallCurveData = [];
     element.data = {
-      'plotData': {'macroValuesByThreshold': macroPrecisionRecallCurveData}
+      'plotData': {'macroValuesByThreshold': {'matrices': macroPrecisionRecallCurveData}}
     };
     element.loading = false;
 
@@ -233,7 +233,7 @@
     ];
     const microPrecisionRecallCurveData = [];
     element.data = {
-      'plotData': {'microValuesByThreshold': microPrecisionRecallCurveData}
+      'plotData': {'microValuesByThreshold': {'matrices': microPrecisionRecallCurveData}}
     };
     element.loading = false;
 
@@ -250,7 +250,7 @@
     ];
     const weightedPrecisionRecallCurveData = [];
     element.data = {
-      'plotData': {'weightedValuesByThreshold': weightedPrecisionRecallCurveData}
+      'plotData': {'weightedValuesByThreshold': {'matrices': weightedPrecisionRecallCurveData}}
     };
     element.loading = false;
 
diff --git a/tensorflow_model_analysis/frontend/tfma-plot/tfma-plot.js b/tensorflow_model_analysis/frontend/tfma-plot/tfma-plot.js
index 4ca56866a8..1120fbc64e 100644
--- a/tensorflow_model_analysis/frontend/tfma-plot/tfma-plot.js
+++ b/tensorflow_model_analysis/frontend/tfma-plot/tfma-plot.js
@@ -194,11 +194,22 @@
      * @return {!Array<!Object>}
      */
     computePrecisionRecallCurveData_: function(data) {
+      return this.getMatricesForPRCurve_(
+          data, tfma.PlotDataFieldNames.PRECISION_RECALL_CURVE_DATA);
+    },
+
+    /**
+     * Extracts the matrices data from the curve plot data with the given key.
+     * @param {?Object} data
+     * @param {string} curveKey
+     * @return {!Array<!Object>}
+     * @private
+     */
+    getMatricesForPRCurve_: function(data, curveKey) {
       const plotData = data && data['plotData'] || {};
-      return plotData[tfma.PlotDataFieldNames.PRECISION_RECALL_CURVE_DATA] &&
-          plotData[tfma.PlotDataFieldNames
-                       .PRECISION_RECALL_CURVE_DATA][tfma.PlotDataFieldNames
-                                                         .CONFUSION_MATRICES] ||
+      const curveData = plotData[curveKey];
+      return curveData &&
+          curveData[tfma.PlotDataFieldNames.CONFUSION_MATRICES] ||
           [];
     },
 
@@ -209,10 +220,8 @@
      * @return {!Array<!Object>}
      */
     computeMacroPrecisionRecallCurveData_: function(data) {
-      const plotData = data && data['plotData'] || {};
-      return plotData[tfma.PlotDataFieldNames
-                          .MACRO_PRECISION_RECALL_CURVE_DATA] ||
-          [];
+      return this.getMatricesForPRCurve_(
+          data, tfma.PlotDataFieldNames.MACRO_PRECISION_RECALL_CURVE_DATA);
     },
 
     /**
@@ -222,10 +231,8 @@
      * @return {!Array<!Object>}
      */
     computeMicroPrecisionRecallCurveData_: function(data) {
-      const plotData = data && data['plotData'] || {};
-      return plotData[tfma.PlotDataFieldNames
-                          .MICRO_PRECISION_RECALL_CURVE_DATA] ||
-          [];
+      return this.getMatricesForPRCurve_(
+          data, tfma.PlotDataFieldNames.MICRO_PRECISION_RECALL_CURVE_DATA);
     },
 
     /**
@@ -235,10 +242,8 @@
      * @return {!Array<!Object>}
      */
     computeWeightedPrecisionRecallCurveData_: function(data) {
-      const plotData = data && data['plotData'] || {};
-      return plotData[tfma.PlotDataFieldNames
-                          .WEIGHTED_PRECISION_RECALL_CURVE_DATA] ||
-          [];
+      return this.getMatricesForPRCurve_(
+          data, tfma.PlotDataFieldNames.WEIGHTED_PRECISION_RECALL_CURVE_DATA);
     },
 
     /**
diff --git a/tensorflow_model_analysis/slicer/BUILD b/tensorflow_model_analysis/slicer/BUILD
deleted file mode 100644
index 50af12d4a1..0000000000
--- a/tensorflow_model_analysis/slicer/BUILD
+++ /dev/null
@@ -1,45 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-py_library(
-    name = "slicer",
-    srcs = ["slicer.py"],
-    visibility = ["//tensorflow_model_analysis:__subpackages__"],
-    deps = [
-        ":slice_accessor",
-        "//tensorflow_model_analysis:expect_six_installed",
-        "//tensorflow_model_analysis:types",
-        "//third_party/py/typing",
-    ],
-)
-
-py_test(
-    name = "slicer_test",
-    srcs = ["slicer_test.py"],
-    deps = [
-        ":slicer",
-        "//tensorflow_model_analysis:expect_six_installed",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//third_party/py/numpy",
-    ],
-)
-
-py_library(
-    name = "slice_accessor",
-    srcs = ["slice_accessor.py"],
-    deps = [
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis:types",
-        "//third_party/py/numpy",
-    ],
-)
-
-py_test(
-    name = "slice_accessor_test",
-    srcs = ["slice_accessor_test.py"],
-    deps = [
-        ":slice_accessor",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/eval_saved_model:encoding",
-        "//third_party/py/numpy",
-    ],
-)
diff --git a/tensorflow_model_analysis/types.py b/tensorflow_model_analysis/types.py
index 644d3f7c85..9d5dd4326c 100644
--- a/tensorflow_model_analysis/types.py
+++ b/tensorflow_model_analysis/types.py
@@ -18,6 +18,7 @@
 
 from __future__ import print_function
 
+import copy
 
 import numpy as np
 import tensorflow as tf
@@ -63,7 +64,35 @@ def is_tensor(obj):
 # example and all its "extractions." Extractions that should be emitted to file.
 # Each Extract has a name, stored as the key of the DictOfExtractedValues.
 DictOfExtractedValues = Dict[Text, Any]
-ExampleAndExtracts = NamedTuple(  # pylint: disable=invalid-name
-    'ExampleAndExtracts',
-    [('example', bytes),
-     ('extracts', DictOfExtractedValues)])
+
+
+class ExampleAndExtracts(
+    NamedTuple('ExampleAndExtracts', [('example', bytes),
+                                      ('extracts', DictOfExtractedValues)])):
+  """Example and extracts."""
+
+  def create_copy_with_shallow_copy_of_extracts(self):
+    """Returns a new copy of this with a shallow copy of extracts.
+
+    This is NOT equivalent to making a shallow copy with copy.copy(this).
+    That does NOT make a shallow copy of the dictionary. An illustration of
+    the differences:
+      a = ExampleAndExtracts(example='content', extracts=dict(apple=[1, 2]))
+
+      # The dictionary is shared (and hence the elements are also shared)
+      b = copy.copy(a)
+      b.extracts['banana'] = 10
+      assert a.extracts['banana'] == 10
+
+      # The dictionary is not shared (but the elements are)
+      c = a.create_copy_with_shallow_copy_of_extracts()
+      c.extracts['cherry'] = 10
+      assert 'cherry' not in a.extracts  # The dictionary is not shared
+      c.extracts['apple'][0] = 100
+      assert a.extracts['apple'][0] == 100  # But the elements are
+
+    Returns:
+      A shallow copy of this object.
+    """
+    return ExampleAndExtracts(
+        example=self.example, extracts=copy.copy(self.extracts))
diff --git a/tensorflow_model_analysis/version.py b/tensorflow_model_analysis/version.py
index 608e8a68d9..83c250300c 100644
--- a/tensorflow_model_analysis/version.py
+++ b/tensorflow_model_analysis/version.py
@@ -15,4 +15,4 @@
 
 # Version string for this release of TFMA.
 # Note that setup.py reads and uses this version.
-VERSION_STRING = '0.9.0'
+VERSION_STRING = '0.9.1'
diff --git a/tensorflow_model_analysis/view/BUILD b/tensorflow_model_analysis/view/BUILD
deleted file mode 100644
index 024d7741f8..0000000000
--- a/tensorflow_model_analysis/view/BUILD
+++ /dev/null
@@ -1,24 +0,0 @@
-licenses(["notice"])  # Apache 2.0
-
-py_library(
-    name = "util",
-    srcs = ["util.py"],
-    deps = [
-        "//tensorflow_model_analysis/api/impl:api_types",
-        "//tensorflow_model_analysis/eval_saved_model/post_export_metrics:metric_keys",
-        "//tensorflow_model_analysis/slicer",
-    ],
-)
-
-py_test(
-    name = "util_test",
-    srcs = ["util_test.py"],
-    deps = [
-        ":util",
-        "//tensorflow_model_analysis:constants",
-        "//tensorflow_model_analysis:expect_tensorflow_installed",
-        "//tensorflow_model_analysis/api/impl:api_types",
-        "//tensorflow_model_analysis/eval_saved_model:testutil",
-        "//tensorflow_model_analysis/slicer",
-    ],
-)